@sixb/ducklake 0.1.8 → 0.2.0

package/src/internal/ducklake-sql-executor.ts

@@ -1,6 +1,13 @@
 import { randomUUID } from "node:crypto"
 import type { DatasetColumnDefinition, DatasetDefinition, DatasetRow } from "@sixb/core"
 import {
+  type AnalyzeSqlTransformInput,
+  type AnalyzeSqlTransformSource,
+  type DatasetColumnLineage,
+  type DatasetColumnRef,
+  type DatasetProducer,
+  type DatasetVersion,
+  type DatasetVersionLineage,
   type DatasetVersionRef,
   type DatasetWriteCommitResult,
   type DatasetWriteMode,
@@ -8,6 +15,10 @@ import {
   type LakeSqlExecutor,
   LakeStorageError,
   type PreviewSqlTransformInput,
+  resolveStepLineage,
+  type SqlTransformAnalysis,
+  type SqlTransformColumnRef,
+  type SqlTransformLineage,
 } from "@sixb/core/lake-storage"
 import type { DuckLakeStorageOptions } from "../types"
 import {
@@ -20,13 +31,17 @@ import { getString } from "./duckdb-row"
 import type { DuckDbQueryRuntime } from "./duckdb-runtime"
 import type { DuckLakeConnectionManager } from "./ducklake-connection-manager"
 import type { DuckLakeDatasetCatalog } from "./ducklake-dataset-catalog"
-import type { DuckLakeSnapshotReader } from "./ducklake-snapshot-reader"
+import type { DuckLakeSnapshotReader, DuckLakeTableVersionRef } from "./ducklake-snapshot-reader"
 import type { DuckLakeWriteCoordinator } from "./ducklake-write-coordinator"
 import { datasetColumnTypeToDuckDbSql } from "./schema"
 import { quoteIdentifier } from "./sql"
+import { analyzeSqlTransformLineage, type SqlLineageAnalysis } from "./sql-lineage/analysis"
 import {
+  createDuckLakeSqlTransformTemplate,
   type DuckLakeSqlTransformSourceRelation,
-  renderDuckLakeSqlTransformSql,
+  type DuckLakeSqlTransformTemplate,
+  pickSourceColumns,
+  renderDuckLakeSourceRelations,
 } from "./sql-transform-relations"
 import { parseVersionId } from "./versions"
 
@@ -35,9 +50,16 @@ const MAX_PREVIEW_LIMIT = 1_000
 
 interface ResolvedDuckLakeSqlTransformSources {
   readonly relations: Readonly<Record<string, DuckLakeSqlTransformSourceRelation>>
+  readonly tables: Readonly<Record<string, ResolvedDuckLakeSourceTable>>
   readonly inputs: readonly DatasetVersionRef[]
 }
 
+/** A source's stored definition, and the version it resolved to. */
+interface ResolvedDuckLakeSourceTable {
+  readonly definition: DatasetDefinition
+  readonly version: DuckLakeTableVersionRef
+}
+
 interface ApplySqlTransformInput {
   readonly runtime: DuckDbQueryRuntime
   readonly target: DatasetDefinition
@@ -79,13 +101,11 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
   async *preview(input: PreviewSqlTransformInput<"duckdb">): AsyncIterable<DatasetRow> {
     this.connections.assertOpen()
 
+    const template = createDuckLakeSqlTransformTemplate(Object.keys(input.sources), input.sql)
+
     const rows = await this.connections.withAttachedRuntime(async (runtime) => {
       const sources = await this.resolveSources(runtime, input.sources)
-      const sql = renderDuckLakeSqlTransformSql({
-        options: this.options,
-        sources: sources.relations,
-        sql: input.sql,
-      })
+      const sql = template.render(renderDuckLakeSourceRelations(this.options, sources.relations))
       return runtime.query(buildPreviewSql(sql, previewLimit(input.limit)))
     })
 
@@ -98,16 +118,29 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
     this.connections.assertOpen()
     assertDatasetWriteMode(input.mode, "SQL transform")
 
+    const template = createDuckLakeSqlTransformTemplate(Object.keys(input.sources), input.sql)
+
     return this.writes.withCommitRuntime(async (runtime) => {
       const target = await this.resolveTargetDataset(runtime, input.target)
       const sources = await this.resolveSources(runtime, input.sources)
-      const sql = renderDuckLakeSqlTransformSql({
-        options: this.options,
-        sources: sources.relations,
-        sql: input.sql,
-      })
+      const sql = template.render(renderDuckLakeSourceRelations(this.options, sources.relations))
 
-      const checkedCasts = await this.assertResultSchemaMatchesTarget(runtime, target, sql)
+      const sourceVersions = await this.readSourceVersions(runtime, sources).catch(
+        (error: unknown) => skipLineage(target, error)
+      )
+
+      // The analysis binds in its own DuckDB instance while the lake checks the result schema.
+      const [checkedCasts, stepLineage] = await Promise.all([
+        this.assertResultSchemaMatchesTarget(runtime, target, sql),
+        sourceVersions === undefined
+          ? undefined
+          : analyzeExecutedLineage(
+              template,
+              pinnedSources(sources, sourceVersions),
+              sources.relations,
+              target
+            ),
+      ])
 
       return this.writes.commitDatasetVersionOnExclusiveRuntime(runtime, {
         dataset: target,
@@ -116,6 +149,10 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
         commitMessage: input.commitMessage ?? `execute SQL transform for dataset ${target.id}`,
         producer: input.producer,
         inputs: sources.inputs,
+        lineage:
+          stepLineage === undefined || sourceVersions === undefined
+            ? unanalyzedLineage(target, input.producer)
+            : resolveStepLineage(stepLineage, [...sourceVersions.values()]),
         applyChanges: (runtime, context) =>
           this.applySqlTransform({
             runtime,
@@ -129,21 +166,47 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
     })
   }
 
+  /**
+   * Column lineage bound against the definitions passed in, without reading the lake: it works
+   * before any source version exists. Their columns are bound in the order given, which can differ
+   * from the lake's stored order; `execute` binds each source's stored version schema instead.
+   */
+  async analyze(input: AnalyzeSqlTransformInput<"duckdb">): Promise<SqlTransformAnalysis> {
+    for (const source of Object.values(input.sources)) {
+      this.datasets.assertSchema(source.dataset)
+    }
+
+    const analysis = await analyzeSqlTransformLineage({
+      sources: input.sources,
+      template: createDuckLakeSqlTransformTemplate(Object.keys(input.sources), input.sql),
+    })
+    return analysis.supported
+      ? { supported: true, lineage: analysis.lineage }
+      : { supported: false, reason: analysis.reason }
+  }
+
   private async resolveSources(
     runtime: DuckDbQueryRuntime,
     sources: PreviewSqlTransformInput<"duckdb">["sources"]
   ): Promise<ResolvedDuckLakeSqlTransformSources> {
     const relations: Record<string, DuckLakeSqlTransformSourceRelation> = Object.create(null)
+    const tables: Record<string, ResolvedDuckLakeSourceTable> = Object.create(null)
     const inputs: DatasetVersionRef[] = []
 
     for (const [sourceName, source] of Object.entries(sources)) {
       const definition = await this.resolveSourceDataset(runtime, source.dataset)
+
+      // Rejects a pick of columns the stored definition does not have before any SQL renders it.
+      pickSourceColumns(sourceName, definition, source.columns)
+
       const versionRef = await this.resolveSourceVersion(runtime, definition, source.versionId)
 
       relations[sourceName] = {
         datasetId: definition.id,
         versionId: versionRef.versionId,
+        ...(source.columns === undefined ? {} : { columns: source.columns }),
       }
+      tables[sourceName] = { definition, version: versionRef }
       inputs.push({
         datasetId: definition.id,
         versionId: versionRef.versionId,
@@ -152,6 +215,7 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
 
     return {
       relations: Object.freeze(relations),
+      tables: Object.freeze(tables),
       inputs: Object.freeze(inputs),
     }
   }
@@ -185,11 +249,25 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
     return definition
   }
 
+  /** The source versions read, by source name, with their schemas from one catalog read. */
+  private async readSourceVersions(
+    runtime: DuckDbQueryRuntime,
+    sources: ResolvedDuckLakeSqlTransformSources
+  ): Promise<ReadonlyMap<string, DatasetVersion>> {
+    const names = Object.keys(sources.tables)
+    const versions = await this.snapshots.getVersionsForRefs(
+      runtime,
+      names.map((sourceName) => sources.tables[sourceName].version)
+    )
+
+    return new Map(names.map((sourceName, index) => [sourceName, versions[index]]))
+  }
+
   private async resolveSourceVersion(
     runtime: DuckDbQueryRuntime,
     dataset: DatasetDefinition,
     versionId?: string
-  ): Promise<DatasetVersionRef> {
+  ): Promise<DuckLakeTableVersionRef> {
     if (versionId !== undefined) {
       const snapshotId = parseVersionId(versionId)
       const versionRef = await this.snapshots.getVersionRefForSnapshot(runtime, dataset, snapshotId)
@@ -301,6 +379,128 @@ export class DuckLakeSqlExecutor implements LakeSqlExecutor<"duckdb"> {
   }
 }
 
+/**
+ * Each source as the transform binds it on the lake: the schema of the version it reads, in the
+ * stored column order, then its pick. Today's definition may have gained columns since, and a new
+ * column can change what a name in the SQL resolves to, such as a lateral alias.
+ */
+function pinnedSources(
+  sources: ResolvedDuckLakeSqlTransformSources,
+  versions: ReadonlyMap<string, DatasetVersion>
+): Readonly<Record<string, AnalyzeSqlTransformSource>> {
+  const pinned: Record<string, AnalyzeSqlTransformSource> = Object.create(null)
+  for (const [sourceName, version] of versions) {
+    pinned[sourceName] = {
+      dataset: { ...sources.tables[sourceName].definition, schema: version.schema },
+      columns: sources.relations[sourceName].columns,
+    }
+  }
+
+  return pinned
+}
+
+/**
+ * Lineage of the SQL about to run, by the sources' dataset ids, against the schemas of the
+ * versions it reads.
+ *
+ * SQL that cannot be analyzed exactly but reads only its sources, such as a PIVOT, depends on at
+ * most every column it reads. Otherwise there is none: the SQL may read what no lineage names. A
+ * result shaped unlike the target means the analysis tables diverged from what runs on the lake, so
+ * that lineage is not recorded either.
+ */
+async function analyzeExecutedLineage(
+  template: DuckLakeSqlTransformTemplate,
+  pinnedSources: Readonly<Record<string, AnalyzeSqlTransformSource>>,
+  relations: Readonly<Record<string, DuckLakeSqlTransformSourceRelation>>,
+  target: DatasetDefinition
+): Promise<DatasetColumnLineage | undefined> {
+  let analysis: SqlLineageAnalysis
+  try {
+    analysis = await analyzeSqlTransformLineage({ sources: pinnedSources, template })
+  } catch (error) {
+    return skipLineage(target, error)
+  }
+
+  if (!analysis.supported) {
+    return analysis.readsOnlySources
+      ? toDatasetColumnLineage(everySourceColumn(pinnedSources, target), relations)
+      : undefined
+  }
+
+  if (!describesTarget(analysis.outputColumns, target)) return undefined
+  return toDatasetColumnLineage(analysis.lineage, relations)
+}
+
+/** Every column the sources expose decides every row of the target, and so each of its values. */
+function everySourceColumn(
+  sources: Readonly<Record<string, AnalyzeSqlTransformSource>>,
+  target: DatasetDefinition
+): SqlTransformLineage {
+  const rows = Object.entries(sources).flatMap(([source, { dataset, columns }]) =>
+    (columns ?? dataset.schema.columns.map((column) => column.name)).map((column) => ({
+      source,
+      column,
+    }))
+  )
+
+  return { rows, columns: Object.fromEntries(target.schema.columns.map(({ name }) => [name, []])) }
+}
+
+/**
+ * A pipeline write without lineage is classified with every marking. Any other write without it
+ * would depend on its own columns only, so one whose SQL cannot be analyzed records that.
+ */
+function unanalyzedLineage(
+  target: DatasetDefinition,
+  producer: DatasetProducer | undefined
+): DatasetVersionLineage | undefined {
+  return producer?.kind === "pipeline"
+    ? undefined
+    : { unknown: `a SQL transform into '${target.id}' could not be analyzed` }
+}
+
+/**
+ * Lineage never fails a transform: invalid SQL or picks fail where the transform binds on the lake.
+ * When the analysis cannot even start, the version records no lineage of its own SQL.
+ */
+function skipLineage(target: DatasetDefinition, error: unknown): undefined {
+  const message = error instanceof Error ? error.message : String(error)
+  console.warn(
+    `[SixbDuckLake] SQL transform into dataset '${target.id}' records no column lineage: ${message}`
+  )
+  return undefined
+}
+
+function describesTarget(outputColumns: readonly string[], target: DatasetDefinition): boolean {
+  return (
+    outputColumns.length === target.schema.columns.length &&
+    target.schema.columns.every((column, index) => outputColumns[index] === column.name)
+  )
+}
+
+/** Source names become dataset ids; two sources reading one dataset share its column refs. */
+function toDatasetColumnLineage(
+  lineage: SqlTransformLineage,
+  relations: Readonly<Record<string, DuckLakeSqlTransformSourceRelation>>
+): DatasetColumnLineage {
+  const toRefs = (refs: readonly SqlTransformColumnRef[]): DatasetColumnRef[] => {
+    const unique = new Map<string, DatasetColumnRef>()
+    for (const { source, column } of refs) {
+      const datasetId = relations[source].datasetId
+      unique.set(JSON.stringify([datasetId, column]), { datasetId, column })
+    }
+
+    return [...unique.values()]
+  }
+
+  return {
+    rows: toRefs(lineage.rows),
+    columns: Object.fromEntries(
+      Object.entries(lineage.columns).map(([column, refs]) => [column, toRefs(refs)])
+    ),
+  }
+}
+
 function previewLimit(limit: number | undefined): number {
   if (limit === undefined || !Number.isFinite(limit)) {
     return DEFAULT_PREVIEW_LIMIT

package/src/internal/ducklake-version-transaction.ts

@@ -0,0 +1,224 @@
+import { LakeConcurrencyError, LakeStorageError } from "@sixb/core/lake-storage"
+import type { DuckLakeCatalogOptions, DuckLakeStorageOptions } from "../types"
+import { getBigIntLike } from "./duckdb-row"
+import type { DuckDbQueryRuntime } from "./duckdb-runtime"
+import { duckLakeAlias, quoteIdentifier } from "./sql"
+
+/**
+ * How a version transaction reruns after other writers moved the catalog under it. Of concurrent
+ * commits to one claimed table, only the first commits: the others rerun, behind a winner that may
+ * already be committing its next version, so rerunning at once can lose every time. Each rerun
+ * after the first waits a random part (full jitter) of a window that doubles up to a cap, which
+ * spreads the losers apart.
+ *
+ * Ten writers that kept committing to one keyed dataset on PostgreSQL needed up to 10 attempts per
+ * write; rerunning at once, one of them ran out of 12. 16 attempts leave room above that, and wait
+ * at most 9.3 s in all (4.6 s on average) before a writer that cannot commit fails.
+ */
+const VERSION_RERUNS = { maxAttempts: 16, baseDelayMillis: 10, maxDelayMillis: 1_000 } as const
+
+/**
+ * The catalog moved while a version transaction was being prepared or committed, before anything
+ * was stored. Rerunning the whole transaction from a fresh read of the latest version is safe.
+ */
+export class CatalogMovedError extends LakeConcurrencyError {}
+
+/** The latest version of a dataset, and the lake's snapshot it was read at. */
+export interface LatestVersionAtSnapshot<T> {
+  readonly snapshotId: string
+  readonly latestVersion: T | null
+}
+
+export interface BeginVersionTransactionInput<T> {
+  readonly options: DuckLakeStorageOptions
+  readonly runtime: DuckDbQueryRuntime
+  readonly datasetId: string
+  readonly signal?: AbortSignal
+  /**
+   * Reads the latest version of the dataset, which the transaction then builds on. Its first
+   * statement reads the lake's current snapshot, which fixes the transaction's snapshot.
+   */
+  readonly readLatestVersion: () => Promise<LatestVersionAtSnapshot<T>>
+}
+
+/**
+ * Begins the transaction of a new dataset version together with a read of the latest version, so
+ * what the new version derives from its parent, such as its row count or its lineage, describes
+ * the parent it commits on top of, provided every version it must not build on unseen claims its
+ * table (see the write coordinator).
+ */
+export async function beginTransactionOnLatestVersion<T>(
+  input: BeginVersionTransactionInput<T>
+): Promise<T | null> {
+  if (!isSqliteCatalog(input.options.catalog)) {
+    await input.runtime.run("BEGIN TRANSACTION")
+
+    // DuckLake fixes a transaction's snapshot on its first read of the lake, not at BEGIN: the
+    // read fixes it before it reads the latest version, so a version that claimed the table and
+    // committed after that read makes this transaction conflict.
+    return (await input.readLatestVersion()).latestVersion
+  }
+
+  // Direct metadata reads enlist a second SQLite transaction on the outer DuckDB connection.
+  // Its read lock blocks DuckLake's own metadata connection at COMMIT. Read in autocommit,
+  // then fix the transaction's DuckLake snapshot and verify nothing changed in between: the
+  // read is then valid for this transaction, including null (no previous version).
+  input.signal?.throwIfAborted()
+  const { snapshotId, latestVersion } = await input.readLatestVersion()
+
+  await input.runtime.run("BEGIN TRANSACTION")
+  if ((await currentSnapshotId(input.options, input.runtime)) !== snapshotId) {
+    await input.runtime.run("ROLLBACK")
+    throw new CatalogMovedError(
+      `[SixbDuckLake] Catalog changed while a write to dataset '${input.datasetId}' was prepared.`
+    )
+  }
+
+  return latestVersion
+}
+
+export interface CommitOnReadSnapshot {
+  /** Recycles the connection when its DuckLake retry setting cannot be restored. */
+  readonly poisonRuntime: () => void
+}
+
+/**
+ * Commits a version transaction. A conflict with a concurrent commit fails before DuckLake stores
+ * anything, and is raised as a `CatalogMovedError` for the transaction to rerun.
+ *
+ * DuckLake commits on top of any snapshot committed since the transaction began that does not
+ * conflict with it. With `onReadSnapshot`, the commit fails on any such snapshot instead, of this
+ * dataset or not: a version that must sit on the version it read, but that inserts and deletes no
+ * rows DuckLake could check, needs it.
+ */
+export async function commitVersionTransaction(
+  runtime: DuckDbQueryRuntime,
+  datasetId: string,
+  onReadSnapshot?: CommitOnReadSnapshot
+): Promise<void> {
+  if (onReadSnapshot) {
+    await runtime.run("SET ducklake_max_retry_count = 0")
+  }
+
+  try {
+    await runtime.run("COMMIT")
+  } catch (error) {
+    throw catalogMovedError(error, datasetId) ?? error
+  } finally {
+    if (onReadSnapshot) {
+      await runtime.run("RESET ducklake_max_retry_count").catch(() => {
+        // The setting is session state on a shared connection, and the commit already succeeded
+        // or failed durably: never change that outcome, but never reuse the connection either.
+        onReadSnapshot.poisonRuntime()
+      })
+    }
+  }
+}
+
+/**
+ * Reruns a version transaction that the catalog moved under, before anything was stored. The
+ * rerun reads the latest version again: one guarded on an expected latest version checks its guard
+ * against what moved, and fails as an optimistic commit if this dataset is what moved.
+ */
+export async function rerunWhenCatalogMoves<T>(
+  datasetId: string,
+  transaction: (attempt: number) => Promise<T>,
+  signal?: AbortSignal
+): Promise<T> {
+  for (let attempt = 1; ; attempt++) {
+    try {
+      return await transaction(attempt)
+    } catch (error) {
+      if (!(error instanceof CatalogMovedError)) {
+        throw error
+      }
+
+      if (attempt === VERSION_RERUNS.maxAttempts) {
+        throw new LakeConcurrencyError(
+          `[SixbDuckLake] Concurrent writers moved the catalog during ${attempt} attempts to commit dataset '${datasetId}'. Retry the write.`,
+          { cause: error }
+        )
+      }
+    }
+
+    await waitBeforeRerun(rerunDelayMillis(attempt), signal)
+  }
+}
+
+/** The wait before the given rerun: none before the first, then full jitter on a doubling window. */
+function rerunDelayMillis(rerun: number): number {
+  if (rerun === 1) {
+    return 0
+  }
+
+  const window = Math.min(
+    VERSION_RERUNS.maxDelayMillis,
+    VERSION_RERUNS.baseDelayMillis * 2 ** (rerun - 1)
+  )
+
+  return Math.random() * window
+}
+
+function waitBeforeRerun(millis: number, signal: AbortSignal | undefined): Promise<void> {
+  signal?.throwIfAborted()
+  if (millis === 0) {
+    return Promise.resolve()
+  }
+
+  return new Promise((resolve, reject) => {
+    const abort = () => {
+      clearTimeout(timer)
+      reject(signal?.reason)
+    }
+    const timer = setTimeout(() => {
+      signal?.removeEventListener("abort", abort)
+      resolve()
+    }, millis)
+
+    signal?.addEventListener("abort", abort, { once: true })
+  })
+}
+
+export function isSqliteCatalog(catalog: DuckLakeCatalogOptions): boolean {
+  return (
+    catalog.type === "sqlite" ||
+    (catalog.type === "custom" && catalog.uri.toLowerCase().startsWith("sqlite:"))
+  )
+}
+
+/**
+ * DuckLake refuses a commit that conflicts with one committed since the transaction began, or that
+ * exhausted its retries against concurrent commits, before it stores anything. Other errors, such
+ * as a lost connection during COMMIT, are not classified: their outcome is ambiguous.
+ */
+function catalogMovedError(error: unknown, datasetId: string): CatalogMovedError | undefined {
+  if (
+    error instanceof Error &&
+    error.message.startsWith(
+      "TransactionContext Error: Failed to commit: Failed to commit DuckLake transaction."
+    ) &&
+    (error.message.includes("Transaction conflict - ") ||
+      error.message.includes("Exceeded the maximum retry count"))
+  ) {
+    return new CatalogMovedError(
+      `[SixbDuckLake] Concurrent catalog commit for dataset '${datasetId}'.`,
+      { cause: error }
+    )
+  }
+
+  return undefined
+}
+
+async function currentSnapshotId(
+  options: DuckLakeStorageOptions,
+  runtime: DuckDbQueryRuntime
+): Promise<string> {
+  const [row] = await runtime.query(
+    `SELECT id FROM ${quoteIdentifier(duckLakeAlias(options))}.current_snapshot()`
+  )
+  if (row === undefined) {
+    throw new LakeStorageError("[SixbDuckLake] DuckLake did not return a current snapshot.")
+  }
+
+  return String(getBigIntLike(row, "id"))
+}