@sixb/ducklake 0.1.6 → 0.1.8

package/src/internal/ducklake-write-coordinator.ts

@@ -18,7 +18,7 @@ import {
   LakeStorageError,
   retryDatasetMergeCommit,
 } from "@sixb/core/lake-storage"
-import type { DuckLakeStorageOptions } from "../types"
+import type { DuckLakeCatalogOptions, DuckLakeStorageOptions } from "../types"
 import { localCatalogCoordinationKey } from "./catalog-key"
 import {
   type ApplyDatasetRowsResult,
@@ -330,13 +330,12 @@ export class DuckLakeWriteCoordinator {
     // which snapshot belongs to this transaction.
     const previousWriteSnapshotId = await this.lastCommittedSnapshotId(input.runtime)
 
-    await input.runtime.run("BEGIN TRANSACTION")
     let committed = false
     try {
       // Guarded Sixb commits need compare-and-swap semantics. DuckLake can
       // automatically retry non-conflicting commits, so guarded transactions
-      // disable retries and re-check the dataset head after BEGIN.
-      const latestVersion = await this.latestVersionForCommit(input)
+      // disable retries and validate a head belonging to the native write snapshot.
+      const latestVersion = await this.beginVersionTransaction(input)
       if (input.expectedLatestVersionId !== undefined) {
         this.assertExpectedLatestVersion({
           datasetId: input.dataset.id,
@@ -527,15 +526,16 @@ export class DuckLakeWriteCoordinator {
   private async ensureNoOpHasSnapshot(
     input: DuckLakeCommitVersionOutcomeRuntimeInput,
     changeResult: ApplyDatasetRowsResult,
-    knownLatestVersion: DuckLakeVersionSummary | null
+    knownLatestVersion: DuckLakeVersionSummary | null | undefined
   ): Promise<void> {
     if (changeResult.dataChangeExpected) {
       return
     }
 
     const latestVersion =
-      knownLatestVersion ??
-      (await this.snapshots.getLatestVersionSummaryForDefinition(input.runtime, input.dataset))
+      knownLatestVersion === undefined
+        ? await this.snapshots.getLatestVersionSummaryForDefinition(input.runtime, input.dataset)
+        : knownLatestVersion
     if (
       latestVersion
         ? input.mode === "merge" || !hasDatasetInputChanges(latestVersion.inputs, input.inputs)
@@ -582,14 +582,49 @@ export class DuckLakeWriteCoordinator {
     }
   }
 
-  private async latestVersionForCommit(
+  private async beginVersionTransaction(
     input: DuckLakeCommitVersionOutcomeRuntimeInput
-  ): Promise<DuckLakeVersionSummary | null> {
-    if (input.expectedLatestVersionId === undefined && input.mode !== "append") {
-      return null
+  ): Promise<DuckLakeVersionSummary | null | undefined> {
+    if (!isSqliteCatalog(this.options.catalog)) {
+      await input.runtime.run("BEGIN TRANSACTION")
+      if (input.expectedLatestVersionId === undefined && input.mode !== "append") {
+        return undefined
+      }
+      return this.snapshots.getLatestVersionSummaryForDefinition(input.runtime, input.dataset)
     }
 
-    return this.snapshots.getLatestVersionSummaryForDefinition(input.runtime, input.dataset)
+    // 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 pin the native DuckLake snapshot and verify nothing changed in between. The summary
+    // is now valid for this write transaction, including null (no previous version).
+    // Only preparation can repeat: staged mutations and COMMIT are never replayed here.
+    for (let attempt = 0; attempt < 3; attempt++) {
+      input.signal?.throwIfAborted()
+      const snapshotId = await this.currentSnapshotId(input.runtime)
+      const latestVersion = await this.snapshots.getLatestVersionSummaryForDefinition(
+        input.runtime,
+        input.dataset
+      )
+      await input.runtime.run("BEGIN TRANSACTION")
+      if ((await this.currentSnapshotId(input.runtime)) === snapshotId) {
+        return latestVersion
+      }
+      await input.runtime.run("ROLLBACK")
+    }
+
+    throw new LakeConcurrencyError(
+      `[SixbDuckLake] Catalog changed during three write preparations for dataset '${input.dataset.id}'. Retry the write.`
+    )
+  }
+
+  private async currentSnapshotId(runtime: DuckDbQueryRuntime): Promise<string> {
+    const [row] = await runtime.query(
+      `SELECT id FROM ${quoteIdentifier(duckLakeAlias(this.options))}.current_snapshot()`
+    )
+    if (row === undefined) {
+      throw new LakeStorageError("[SixbDuckLake] DuckLake did not return a current snapshot.")
+    }
+    return String(getBigIntLike(row, "id"))
   }
 
   private async lastCommittedSnapshotId(runtime: DuckDbQueryRuntime): Promise<string | null> {
@@ -681,15 +716,21 @@ export class DuckLakeWriteCoordinator {
   }
 }
 
+function isSqliteCatalog(catalog: DuckLakeCatalogOptions): boolean {
+  return (
+    catalog.type === "sqlite" ||
+    (catalog.type === "custom" && catalog.uri.toLowerCase().startsWith("sqlite:"))
+  )
+}
+
 function commitRowCount(
   input: Pick<DuckLakeCommitVersionOutcomeInput, "expectedLatestVersionId" | "mode">,
   changeResult: ApplyDatasetRowsResult,
-  latestVersion: DuckLakeVersionSummary | null
+  latestVersion: DuckLakeVersionSummary | null | undefined
 ): CommitRowCount {
   // Unguarded append commits can be retried by DuckLake after this transaction
   // read the previous latest version. In that case previousRowCount +
-  // sourceRowCount can underreport the committed snapshot, so omit metadata
-  // and let hydration count the exact committed snapshot instead.
+  // sourceRowCount can underreport the committed snapshot, so leave the count unknown.
   if (input.mode === "append" && input.expectedLatestVersionId === undefined) {
     return { kind: "unknown" }
   }

package/src/internal/ducklake-write-session.ts

@@ -183,6 +183,7 @@ class DuckLakeWriteSession implements LakeWriteSession {
     }
 
     this.cleanedUp = true
+    this.stagedPrimaryKeys.clear()
     await this.dropStagingTable()
   }
 

package/src/internal/sql.ts

@@ -60,21 +60,50 @@ export function buildSetPostgresPoolSql(
  * Render a fully-qualified DuckLake metadata table path.
  *
  * PostgreSQL catalogs expose DuckLake metadata tables under the configured
- * metadata schema inside the attached metadata catalog. Local catalogs use
- * DuckDB's default `main` schema.
+ * metadata schema (PostgreSQL's `public` by default) inside the attached
+ * metadata catalog. Local catalogs use DuckDB's default `main` schema.
  */
 export function duckLakeMetadataTableName(
   options: Pick<DuckLakeStorageOptions, "alias" | "catalog">,
   tableName: string
 ): string {
-  const schemaName =
-    options.catalog.type === "postgres" ? (options.catalog.metadataSchema ?? "main") : "main"
-
   return `${quoteIdentifier(duckLakeMetadataCatalog(options))}.${quoteIdentifier(
-    schemaName
+    duckLakeMetadataSchema(options)
   )}.${quoteIdentifier(tableName)}`
 }
 
+/**
+ * Render a metadata query that runs where the catalog lives.
+ *
+ * DuckDB's PostgreSQL scanner pushes comparisons and constant IN lists down to PostgreSQL, but not
+ * string matching, so a query that filters on text would copy the whole column into DuckDB. On a
+ * PostgreSQL catalog the query therefore runs server-side through `postgres_query`; elsewhere the
+ * catalog is local and DuckDB runs it. `build` must produce SQL that both dialects accept.
+ */
+export function buildDuckLakeMetadataQuery(
+  options: Pick<DuckLakeStorageOptions, "alias" | "catalog">,
+  build: (table: (tableName: string) => string) => string
+): string {
+  if (options.catalog.type !== "postgres") {
+    return build((tableName) => duckLakeMetadataTableName(options, tableName))
+  }
+
+  // Unqualified names follow the connection's search path, as DuckLake's own writes do.
+  const schema = options.catalog.metadataSchema
+  const sql = build((tableName) =>
+    schema === undefined
+      ? quoteIdentifier(tableName)
+      : `${quoteIdentifier(schema)}.${quoteIdentifier(tableName)}`
+  )
+  return `SELECT * FROM postgres_query(${quoteSqlString(duckLakeMetadataCatalog(options))}, ${quoteSqlString(sql)})`
+}
+
+function duckLakeMetadataSchema(options: Pick<DuckLakeStorageOptions, "catalog">): string {
+  if (options.catalog.type !== "postgres") return "main"
+  // Without a metadata schema DuckLake writes to PostgreSQL's default schema.
+  return options.catalog.metadataSchema ?? "public"
+}
+
 /**
  * Quote a SQL identifier after validating it is already identifier-shaped.
  *

package/src/types.ts

@@ -151,6 +151,8 @@ export interface DuckLakeStorageOptions {
   readonly alias?: string
   readonly duckdb?: DuckDbRuntimeOptions
   readonly postgresPool?: DuckLakePostgresPoolOptions
+  /** Maximum native streaming readers per provider; extra reads use pages. Default: 4. */
+  readonly maxStreamingReads?: number
   readonly secrets?: readonly DuckDbSecretOptions[]
   /** Runs after DuckLake/catalog extensions load and before ATTACH. */
   readonly setupSql?: readonly string[]