@sixb/ducklake 0.1.8 → 0.1.9

package/src/internal/sql-lineage/analysis.ts

@@ -0,0 +1,257 @@
+import { randomUUID } from "node:crypto"
+import { type DuckDBConnection, DuckDBInstance } from "@duckdb/node-api"
+import type { DatasetColumnDefinition } from "@sixb/core"
+import type {
+  AnalyzeSqlTransformSource,
+  SqlTransformColumnRef,
+  SqlTransformLineage,
+} from "@sixb/core/lake-storage"
+import { datasetColumnToDuckDbSql } from "../schema"
+import { quoteIdentifier, quoteSqlString } from "../sql"
+import {
+  type DuckLakeSqlTransformTemplate,
+  pickSourceColumns,
+  renderSourceRelation,
+} from "../sql-transform-relations"
+import { parseSerializerOutput } from "./plan-json"
+import { type PlanSourceTables, planLineage } from "./plan-walker"
+import { rejectPlanFormat } from "./rejection"
+import { assertSingleSelect, readsOnlySources } from "./statement-gate"
+
+interface AnalyzeSqlTransformLineageInput {
+  readonly sources: Readonly<Record<string, AnalyzeSqlTransformSource>>
+  readonly template: DuckLakeSqlTransformTemplate
+}
+
+/**
+ * `SqlTransformAnalysis` with the result's column names in order, which the lineage record cannot
+ * carry: object keys that look like integers enumerate first.
+ */
+export type SqlLineageAnalysis =
+  | {
+      readonly supported: true
+      readonly lineage: SqlTransformLineage
+      readonly outputColumns: readonly string[]
+    }
+  | { readonly supported: false; readonly reason: string; readonly readsOnlySources: boolean }
+
+interface AnalyzedSql {
+  readonly lineage: SqlTransformLineage
+  readonly outputColumns: readonly string[]
+}
+
+/** An empty table standing in for one source, under a name the transform cannot guess. */
+interface AnalysisTable {
+  readonly sourceName: string
+  readonly tableName: string
+  readonly columns: readonly DatasetColumnDefinition[]
+  readonly relationSql: string
+}
+
+/**
+ * No file, network, extension, or setting is reachable from the analysis instance; the transform
+ * only binds against its sources' empty tables.
+ */
+const ISOLATED_INSTANCE_CONFIG = {
+  enable_external_access: "false",
+  autoinstall_known_extensions: "false",
+  autoload_known_extensions: "false",
+}
+
+/**
+ * Column lineage of a SQL transform, derived without reading any data.
+ *
+ * A fresh in-memory DuckDB instance creates one empty table per source from its column schema,
+ * the transform renders against those tables in the same relation shape it runs with on the lake,
+ * and DuckDB binds it. The bound, unoptimized plan is then walked operator by operator. Whatever
+ * the walker does not understand exactly is rejected rather than guessed; a rejection or a bind
+ * error is an unsupported analysis, while an invalid source pick throws. An unsupported analysis
+ * of SQL that parsed still says whether its parse tree reads anything but the sources.
+ */
+export async function analyzeSqlTransformLineage(
+  input: AnalyzeSqlTransformLineageInput
+): Promise<SqlLineageAnalysis> {
+  const tables = analysisTables(input.sources)
+  const sql = input.template.render(
+    new Map(tables.map((table) => [table.sourceName, table.relationSql]))
+  )
+
+  try {
+    return await withAnalysisConnection(tables, async (connection) => {
+      const statements = await serialize(
+        connection,
+        `json_serialize_sql(${quoteSqlString(sql)}, skip_null := false, skip_empty := false)`,
+        "statements",
+        "DuckDB could not parse the transform"
+      )
+
+      try {
+        return { supported: true, ...(await analyzeSql(connection, sql, statements, tables)) }
+      } catch (error) {
+        const sourceTables = new Set(tables.map((table) => table.tableName))
+        const functions = await builtinFunctions(connection).catch(() => undefined)
+        return unsupported(
+          error,
+          functions !== undefined && readsOnlySources(statements, sourceTables, functions)
+        )
+      }
+    })
+  } catch (error) {
+    return unsupported(error, false)
+  }
+}
+
+/** The functions DuckDB defines, the same in every fresh instance: read from the first one. */
+let builtins: Promise<ReadonlySet<string>> | undefined
+
+function builtinFunctions(connection: DuckDBConnection): Promise<ReadonlySet<string>> {
+  if (builtins) return builtins
+
+  const read = connection
+    .runAndReadAll("SELECT DISTINCT lower(function_name) FROM duckdb_functions()")
+    .then((reader) => new Set(reader.getRows().map(([name]) => String(name))))
+
+  // A failed read is not kept: the next analysis reads again.
+  builtins = read
+  read.catch(() => {
+    if (builtins === read) builtins = undefined
+  })
+  return read
+}
+
+function unsupported(error: unknown, readsOnlySources: boolean): SqlLineageAnalysis {
+  return {
+    supported: false,
+    reason: error instanceof Error ? error.message : String(error),
+    readsOnlySources,
+  }
+}
+
+function analysisTables(
+  sources: Readonly<Record<string, AnalyzeSqlTransformSource>>
+): readonly AnalysisTable[] {
+  const token = randomUUID().replaceAll("-", "")
+
+  return Object.entries(sources).map(([sourceName, source], index) => {
+    const tableName = `sixb_lineage_${token}_${index}`
+    return {
+      sourceName,
+      tableName,
+      columns: pickSourceColumns(sourceName, source.dataset, source.columns),
+      relationSql: renderSourceRelation(quoteIdentifier(tableName), source.columns),
+    }
+  })
+}
+
+async function withAnalysisConnection<T>(
+  tables: readonly AnalysisTable[],
+  use: (connection: DuckDBConnection) => Promise<T>
+): Promise<T> {
+  const instance = await DuckDBInstance.create(":memory:", ISOLATED_INSTANCE_CONFIG)
+  try {
+    const connection = await instance.connect()
+    try {
+      for (const table of tables) {
+        const columnsSql = table.columns.map((column) => datasetColumnToDuckDbSql(column))
+        await connection.run(
+          `CREATE TABLE ${quoteIdentifier(table.tableName)} (${columnsSql.join(", ")})`
+        )
+      }
+
+      // The transform cannot loosen the restrictions above through SET or a setting function.
+      await connection.run("SET lock_configuration = true")
+      return await use(connection)
+    } finally {
+      connection.closeSync()
+    }
+  } finally {
+    instance.closeSync()
+  }
+}
+
+async function analyzeSql(
+  connection: DuckDBConnection,
+  sql: string,
+  statements: readonly unknown[],
+  tables: readonly AnalysisTable[]
+): Promise<AnalyzedSql> {
+  assertSingleSelect(statements)
+
+  const plans = await serialize(
+    connection,
+    `json_serialize_plan(${quoteSqlString(sql)}, skip_null := false, skip_empty := false, optimize := false)`,
+    "plans",
+    "DuckDB could not bind the transform"
+  )
+  if (plans.length !== 1) {
+    rejectPlanFormat(`expected one plan, got ${plans.length}`)
+  }
+
+  const sourceColumns = numberSourceColumns(tables)
+  const lineage = planLineage(plans[0], sourceColumns.ids)
+  const names = await outputColumnNames(connection, sql)
+  if (lineage.columns.length !== names.length) {
+    rejectPlanFormat(`the plan has ${lineage.columns.length} columns, the result ${names.length}`)
+  }
+
+  const refs = (ids: readonly number[]): SqlTransformColumnRef[] =>
+    ids.map((id) => sourceColumns.refs[id])
+
+  return {
+    lineage: {
+      rows: refs(lineage.rows),
+      columns: Object.fromEntries(names.map((name, index) => [name, refs(lineage.columns[index])])),
+    },
+    outputColumns: names,
+  }
+}
+
+/** Numbers every column of every table, in source then column order, for the plan walker. */
+function numberSourceColumns(tables: readonly AnalysisTable[]): {
+  readonly ids: PlanSourceTables
+  readonly refs: readonly SqlTransformColumnRef[]
+} {
+  const refs: SqlTransformColumnRef[] = []
+  const ids = new Map(
+    tables.map((table) => {
+      const columnIds = table.columns.map((column) => {
+        refs.push({ source: table.sourceName, column: column.name })
+        return [column.name, refs.length - 1] as const
+      })
+      return [table.tableName, new Map(columnIds)] as const
+    })
+  )
+
+  return { ids, refs }
+}
+
+/**
+ * Names as the executor materializes them: selected through a subquery, where DuckDB renames
+ * duplicate output names (`id`, `id_1`).
+ */
+async function outputColumnNames(
+  connection: DuckDBConnection,
+  sql: string
+): Promise<readonly string[]> {
+  const prepared = await connection.prepare(`SELECT * FROM (${sql}) AS sixb_lineage_result`)
+  try {
+    return Array.from({ length: prepared.columnCount }, (_, index) => prepared.columnName(index))
+  } finally {
+    prepared.destroySync()
+  }
+}
+
+async function serialize(
+  connection: DuckDBConnection,
+  serializerSql: string,
+  listKey: string,
+  failure: string
+): Promise<readonly unknown[]> {
+  const reader = await connection.runAndReadAll(`SELECT ${serializerSql}`)
+  const text = reader.getRows()[0]?.[0]
+  if (typeof text !== "string") {
+    rejectPlanFormat("the serializer returned no JSON")
+  }
+
+  return parseSerializerOutput(text, listKey, failure)
+}

package/src/internal/sql-lineage/expressions.ts

@@ -0,0 +1,208 @@
+import {
+  asObject,
+  booleanField,
+  field,
+  indexField,
+  isJsonObject,
+  type JsonObject,
+  listField,
+  nullableObjectField,
+  objectField,
+  objectListField,
+  stringField,
+} from "./plan-json"
+import { addAll, bindingKey, itemAt, type PlanColumn } from "./plan-relation"
+import { reject, rejectPlanFormat } from "./rejection"
+
+/** Source columns whose values flow into an expression evaluated over an operator's `input`. */
+export function expressionDependencies(
+  expression: unknown,
+  input: readonly PlanColumn[]
+): Set<number> {
+  const dependencies = new Set<number>()
+  addExpression(dependencies, expression, input)
+  return dependencies
+}
+
+/** The sort expressions of an ORDER BY list (`{ type, null_order, expression }` entries). */
+export function orderExpressions(orders: readonly unknown[]): readonly unknown[] {
+  return orders.map((order) => field(asObject(order, "sort order"), "expression"))
+}
+
+/** `{ orders: [...] }` or null, as aggregates and DISTINCT ON write their ORDER BY. */
+export function nestedOrderExpressions(object: JsonObject, key: string): readonly unknown[] {
+  const orderBy = nullableObjectField(object, key)
+  return orderBy === null ? [] : orderExpressions(listField(orderBy, "orders"))
+}
+
+function addExpression(target: Set<number>, value: unknown, input: readonly PlanColumn[]): void {
+  const expression = asObject(value, "expression")
+  const expressionClass = stringField(expression, "expression_class")
+
+  switch (expressionClass) {
+    case "BOUND_REF":
+      // Positional: the index-th column of the operator's input.
+      addAll(target, itemAt(input, indexField(expression, "index"), "column").dependencies)
+      return
+
+    case "BOUND_COLUMN_REF":
+      addAll(target, columnByBinding(expression, input).dependencies)
+      return
+
+    case "BOUND_CONSTANT":
+      return
+
+    case "BOUND_CAST":
+    case "BOUND_UNNEST":
+      addExpression(target, field(expression, "child"), input)
+      return
+
+    case "BOUND_FUNCTION":
+      // A lambda body lives in the bind data with its own column numbering, and every column it
+      // captures is also passed as an extra child.
+      assertNoHiddenExpressions(expression, "lambda_expr")
+      addExpressions(target, listField(expression, "children"), input)
+      return
+
+    case "BOUND_AGGREGATE":
+      // FILTER and ORDER BY shape the aggregated value, as a CASE inside the aggregate would.
+      assertNoHiddenExpressions(expression)
+      addExpressions(target, listField(expression, "children"), input)
+      addOptionalExpression(target, field(expression, "filter"), input)
+      addExpressions(target, nestedOrderExpressions(expression, "order_bys"), input)
+      return
+
+    case "BOUND_WINDOW":
+      addWindowExpression(target, expression, input)
+      return
+
+    case "BOUND_COMPARISON":
+      addExpression(target, field(expression, "left"), input)
+      addExpression(target, field(expression, "right"), input)
+      return
+
+    case "BOUND_BETWEEN":
+      addExpression(target, field(expression, "input"), input)
+      addExpression(target, field(expression, "lower"), input)
+      addExpression(target, field(expression, "upper"), input)
+      return
+
+    case "BOUND_CONJUNCTION":
+    case "BOUND_OPERATOR":
+      addExpressions(target, listField(expression, "children"), input)
+      return
+
+    case "BOUND_CASE":
+      for (const check of objectListField(expression, "case_checks")) {
+        addExpression(target, field(check, "when_expr"), input)
+        addExpression(target, field(check, "then_expr"), input)
+      }
+
+      addExpression(target, field(expression, "else_expr"), input)
+      return
+
+    default:
+      // Subqueries are always flattened into joins by the planner; parameters, lambda
+      // references outside a lambda body, and anything newer are not understood.
+      reject(`expression '${expressionClass}' is not supported`)
+  }
+}
+
+/**
+ * PARTITION BY and ORDER BY are value dependencies of the window column; the window operator also
+ * adds them to the row dependencies. FILTER, frame bounds, offsets and defaults shape the value.
+ */
+function addWindowExpression(
+  target: Set<number>,
+  expression: JsonObject,
+  input: readonly PlanColumn[]
+): void {
+  // Only window aggregates serialize a function, with its bind data; ranking and value functions
+  // such as row_number() or lag() have none.
+  if (stringField(expression, "type") === "WINDOW_AGGREGATE") {
+    assertNoHiddenExpressions(expression)
+  }
+
+  addExpressions(target, listField(expression, "children"), input)
+  addExpressions(target, listField(expression, "partitions"), input)
+  addExpressions(target, orderExpressions(listField(expression, "orders")), input)
+  addExpressions(target, orderExpressions(listField(expression, "arg_orders")), input)
+
+  for (const key of ["filters", "start_expr", "end_expr", "offset_expr", "default_expr"]) {
+    addOptionalExpression(target, field(expression, key), input)
+  }
+}
+
+function addExpressions(
+  target: Set<number>,
+  expressions: readonly unknown[],
+  input: readonly PlanColumn[]
+): void {
+  for (const expression of expressions) {
+    addExpression(target, expression, input)
+  }
+}
+
+function addOptionalExpression(
+  target: Set<number>,
+  expression: unknown,
+  input: readonly PlanColumn[]
+): void {
+  if (expression !== null) {
+    addExpression(target, expression, input)
+  }
+}
+
+/**
+ * DuckDB leaves the subtree under a table function's input unresolved: columns are referenced by
+ * their plan binding instead of their position.
+ */
+function columnByBinding(expression: JsonObject, input: readonly PlanColumn[]): PlanColumn {
+  if (indexField(expression, "depth") !== 0) {
+    reject("correlated column references that the planner did not flatten are not supported")
+  }
+
+  const binding = objectField(expression, "binding")
+  const key = bindingKey(indexField(binding, "table_index"), indexField(binding, "column_index"))
+  const column = input.find((candidate) => candidate.binding === key)
+  if (column === undefined) {
+    rejectPlanFormat(`column binding ${key} does not resolve`)
+  }
+
+  return column
+}
+
+/**
+ * Function bind data is opaque to the analysis. It may hold the lambda body of a list function,
+ * whose captures are also children; any other expression in it would be a hidden dependency.
+ */
+function assertNoHiddenExpressions(expression: JsonObject, allowedKey?: string): void {
+  // Functions without a serializer rebind from their children: there is no bind data to read.
+  if (!booleanField(expression, "has_serialize")) {
+    return
+  }
+
+  const bindData = field(expression, "function_data")
+  const inspected =
+    isJsonObject(bindData) && allowedKey !== undefined
+      ? Object.entries(bindData).flatMap(([key, value]) => (key === allowedKey ? [] : [value]))
+      : [bindData]
+
+  if (inspected.some(containsExpression)) {
+    reject(
+      `function '${stringField(expression, "name")}' keeps expressions in its bind data, which is not supported`
+    )
+  }
+}
+
+function containsExpression(value: unknown): boolean {
+  if (Array.isArray(value)) {
+    return value.some((item) => containsExpression(item))
+  }
+
+  if (!isJsonObject(value)) {
+    return false
+  }
+
+  return Object.hasOwn(value, "expression_class") || Object.values(value).some(containsExpression)
+}

package/src/internal/sql-lineage/plan-json.ts

@@ -0,0 +1,130 @@
+import { reject, rejectPlanFormat } from "./rejection"
+
+/**
+ * Strict readers for DuckDB's JSON serializers (`json_serialize_sql`, `json_serialize_plan`).
+ *
+ * The format is not a stable DuckDB API. Every field the analysis reads must be present with the
+ * expected shape, so a renamed or retyped field in a future release rejects the transform instead
+ * of silently dropping a dependency.
+ */
+export type JsonObject = Readonly<Record<string, unknown>>
+
+export function isJsonObject(value: unknown): value is JsonObject {
+  return typeof value === "object" && value !== null && !Array.isArray(value)
+}
+
+export function asObject(value: unknown, what: string): JsonObject {
+  if (!isJsonObject(value)) {
+    rejectPlanFormat(`${what} is not an object`)
+  }
+
+  return value
+}
+
+/** A field the serializer always writes. It may be null. */
+export function field(object: JsonObject, key: string): unknown {
+  if (!Object.hasOwn(object, key)) {
+    rejectPlanFormat(`missing field '${key}'`)
+  }
+
+  return object[key]
+}
+
+export function objectField(object: JsonObject, key: string): JsonObject {
+  return asObject(field(object, key), `field '${key}'`)
+}
+
+export function nullableObjectField(object: JsonObject, key: string): JsonObject | null {
+  const value = field(object, key)
+  return value === null ? null : asObject(value, `field '${key}'`)
+}
+
+export function stringField(object: JsonObject, key: string): string {
+  const value = field(object, key)
+  if (typeof value !== "string") {
+    rejectPlanFormat(`field '${key}' is not a string`)
+  }
+
+  return value
+}
+
+export function booleanField(object: JsonObject, key: string): boolean {
+  const value = field(object, key)
+  if (typeof value !== "boolean") {
+    rejectPlanFormat(`field '${key}' is not a boolean`)
+  }
+
+  return value
+}
+
+/**
+ * Indexes and table indexes. DuckDB writes 2^64-1 as an "unset" sentinel, which `JSON.parse`
+ * rounds; it stays an integer far past any column count, which is all the analysis compares.
+ */
+export function indexField(object: JsonObject, key: string): number {
+  return asIndex(field(object, key), `field '${key}'`)
+}
+
+function asIndex(value: unknown, what: string): number {
+  if (typeof value !== "number" || !Number.isInteger(value) || value < 0) {
+    rejectPlanFormat(`${what} is not an index`)
+  }
+
+  return value
+}
+
+/** A list field. The serializer writes an absent list as null or as an empty array. */
+export function listField(object: JsonObject, key: string): readonly unknown[] {
+  return asList(field(object, key), `field '${key}'`)
+}
+
+export function asList(value: unknown, what: string): readonly unknown[] {
+  if (value === null) {
+    return []
+  }
+
+  if (!Array.isArray(value)) {
+    rejectPlanFormat(`${what} is not a list`)
+  }
+
+  return value
+}
+
+export function objectListField(object: JsonObject, key: string): readonly JsonObject[] {
+  return listField(object, key).map((item) => asObject(item, `an item of '${key}'`))
+}
+
+export function indexListField(object: JsonObject, key: string): readonly number[] {
+  return asIndexList(field(object, key), `field '${key}'`)
+}
+
+export function asIndexList(value: unknown, what: string): readonly number[] {
+  return asList(value, what).map((item) => asIndex(item, `an item of ${what}`))
+}
+
+export function stringListField(object: JsonObject, key: string): readonly string[] {
+  return listField(object, key).map((item) => {
+    if (typeof item !== "string") {
+      rejectPlanFormat(`an item of '${key}' is not a string`)
+    }
+
+    return item
+  })
+}
+
+/**
+ * Both serializers answer `{ "error": false, <listKey>: [...] }`, or `{ "error": true,
+ * "error_message": ... }` when DuckDB could not parse or bind the query.
+ */
+export function parseSerializerOutput(
+  text: string,
+  listKey: string,
+  failure: string
+): readonly unknown[] {
+  const output = asObject(JSON.parse(text), "serializer output")
+  if (booleanField(output, "error")) {
+    reject(`${failure}: ${stringField(output, "error_message")}`)
+  }
+
+  return listField(output, listKey)
+}

package/src/internal/sql-lineage/plan-relation.ts

@@ -0,0 +1,59 @@
+import { rejectPlanFormat } from "./rejection"
+
+/** Source columns, as ids assigned by the analysis session, that something depends on. */
+export type Dependencies = ReadonlySet<number>
+
+/**
+ * A planner-internal column past the end of a table's columns: `rowid`, or the placeholder a
+ * `count(*)` scan reads. Scanning one is allowed; depending on one rejects the transform.
+ */
+export const VIRTUAL_COLUMN = -1
+
+/** A column flowing out of a plan operator. */
+export interface PlanColumn {
+  readonly dependencies: Dependencies
+  /**
+   * `table_index.column_index` in the plan. Most expressions reference their input by position
+   * (BOUND_REF); the subtree under a table function's input keeps binding references instead.
+   */
+  readonly binding?: string
+}
+
+/** The output of a plan operator: its columns, and what decides which rows exist and their order. */
+export interface PlanRelation {
+  readonly columns: readonly PlanColumn[]
+  readonly rows: Dependencies
+}
+
+export function bindingKey(tableIndex: number, columnIndex: number): string {
+  return `${tableIndex}.${columnIndex}`
+}
+
+/** Positional access into an operator's columns or index lists. */
+export function itemAt<T>(items: readonly T[], index: number, what: string): T {
+  const item = items[index]
+  if (item === undefined) {
+    rejectPlanFormat(`${what} ${index} is out of range`)
+  }
+
+  return item
+}
+
+export function unionOf(...sets: readonly Dependencies[]): Set<number> {
+  const union = new Set<number>()
+  for (const set of sets) {
+    addAll(union, set)
+  }
+
+  return union
+}
+
+export function addAll(target: Set<number>, source: Iterable<number>): void {
+  for (const item of source) {
+    target.add(item)
+  }
+}
+
+export function sameDependencies(left: Dependencies, right: Dependencies): boolean {
+  return left.size === right.size && [...left].every((item) => right.has(item))
+}