@sixb/pg 0.1.9 → 0.1.11

package/src/ontology-storage/source-roots.ts

@@ -0,0 +1,251 @@
+import type { ProjectionEntityRef } from "@sixb/core/internal/materialization"
+import {
+  MaterializationConflictError,
+  MaterializationValidationError,
+  projectionEntityKey,
+} from "@sixb/core/internal/materialization"
+import { sourceStageRoots, utf8SortKey } from "@sixb/core/internal/ontology-storage-provider"
+import type { SourceActivationWrite, StageSourceRowsInput } from "@sixb/core/storage"
+import type { SQLClient } from "../pg-client"
+import type { PgOntologySourceAssertionRow } from "./shared"
+import { jsonParameter, type PgOntologySourceRow } from "./shared"
+
+export function replacementSourceRows(
+  sql: SQLClient,
+  input: {
+    readonly projectId: string
+    readonly sourceId: string
+    readonly materializationId: string
+    readonly incremental: boolean
+  }
+) {
+  const roots = input.incremental
+    ? sql`
+    SELECT roots.* FROM ontology_source_roots AS changed
+    CROSS JOIN LATERAL (
+      SELECT roots.* FROM ontology_source_roots AS roots
+      WHERE roots.project_id = changed.project_id AND roots.source_id = changed.source_id
+        AND roots.root_sort_key = changed.root_sort_key AND roots.active OFFSET 0
+    ) AS roots
+    WHERE changed.project_id = ${input.projectId} AND changed.source_id = ${input.sourceId}
+      AND changed.materialization_id = ${input.materializationId}
+  `
+    : sql`SELECT * FROM ontology_source_roots WHERE project_id = ${input.projectId} AND source_id = ${input.sourceId} AND active`
+  return sql`
+    SELECT * FROM ontology_source_rows WHERE project_id = ${input.projectId}
+      AND source_id = ${input.sourceId} AND materialization_id = ${input.materializationId}
+    UNION ALL
+    SELECT rows.* FROM (${roots}) AS roots
+    CROSS JOIN LATERAL (
+      SELECT rows.* FROM ontology_source_rows AS rows WHERE rows.project_id = roots.project_id
+        AND rows.source_id = roots.source_id AND rows.materialization_id = roots.materialization_id
+        AND rows.root_sort_key = roots.root_sort_key OFFSET 0
+    ) AS rows
+  `
+}
+
+export async function stageSourceRoots(
+  sql: SQLClient,
+  manifest: PgOntologySourceRow,
+  input: StageSourceRowsInput
+): Promise<void> {
+  const roots = sourceStageRoots(manifest.projection_kind, input)
+  if (roots.length === 0) return
+  if (!manifest.base_materialization_id && roots.some((root) => root.deleted)) {
+    throw new MaterializationValidationError("Root deletions require a source delta base.")
+  }
+  const values = jsonParameter(
+    sql,
+    roots.map((root) => ({
+      root_sort_key: root.sortKey,
+      root_kind: root.root.kind,
+      root_key: JSON.parse(root.rootKey),
+      root: root.root,
+      staging_ordinal: root.stagingOrdinal,
+      deleted: root.deleted,
+    }))
+  )
+  const inserted = await sql<{ root_sort_key: string }[]>`
+    INSERT INTO ontology_source_roots (
+      project_id, source_id, materialization_id, root_sort_key, root_kind, root_key, root, staging_ordinal, deleted
+    ) SELECT ${input.projectId}, ${input.source.projectionId}, ${input.materializationId},
+      root_sort_key, root_kind, root_key, root, staging_ordinal, deleted
+    FROM jsonb_to_recordset(${values}) AS input(
+      root_sort_key TEXT, root_kind TEXT, root_key JSONB, root JSONB, staging_ordinal BIGINT, deleted BOOLEAN)
+    ON CONFLICT DO NOTHING RETURNING root_sort_key
+  `
+  if (inserted.length === roots.length) return
+  const insertedKeys = new Set(inserted.map((row) => row.root_sort_key))
+  const skipped = roots.filter((root) => !insertedKeys.has(root.sortKey))
+  const stored = await sql<
+    { root_sort_key: string; staging_ordinal: string | number; deleted: boolean }[]
+  >`
+    SELECT root_sort_key, staging_ordinal, deleted FROM ontology_source_roots
+    WHERE project_id = ${input.projectId} AND source_id = ${input.source.projectionId}
+      AND materialization_id = ${input.materializationId}
+      AND root_sort_key = ANY(${sql.array(skipped.map((root) => root.sortKey))}::text[])
+  `
+  const existing = new Map(stored.map((root) => [root.root_sort_key, root]))
+  if (
+    skipped.some((root) => {
+      const previous = existing.get(root.sortKey)
+      return (
+        !previous ||
+        Number(previous.staging_ordinal) !== root.stagingOrdinal ||
+        previous.deleted !== root.deleted
+      )
+    })
+  )
+    throw new MaterializationValidationError(
+      "Source materialization repeats root or stream ordinal with different content."
+    )
+}
+
+export async function activateSourceRoots(
+  sql: SQLClient,
+  projectId: string,
+  candidate: PgOntologySourceRow,
+  activation: SourceActivationWrite
+): Promise<void> {
+  if (
+    candidate.base_materialization_id !== null &&
+    (candidate.base_materialization_id !== activation.expected.activeMaterializationId ||
+      candidate.base_commit_id !== activation.expected.lastCommitId)
+  )
+    throw new MaterializationConflictError(
+      "projection-fence",
+      "Source delta base changed before activation."
+    )
+  if (candidate.base_materialization_id !== null && Number(candidate.root_count) === 0) return
+  const affected =
+    candidate.base_materialization_id === null
+      ? sql``
+      : sql`AND root_sort_key IN (
+    SELECT root_sort_key FROM ontology_source_roots
+    WHERE project_id = ${projectId} AND source_id = ${activation.source.projectionId}
+      AND materialization_id = ${activation.materializationId}
+  )`
+  await sql`
+    UPDATE ontology_source_roots AS roots SET active = FALSE, retired_at = ${activation.updatedAt}
+    WHERE project_id = ${projectId} AND source_id = ${activation.source.projectionId} AND active
+    ${affected}
+  `
+  await sql`
+    UPDATE ontology_source_roots SET active = NOT deleted,
+      retired_at = CASE WHEN deleted THEN ${activation.updatedAt}::timestamptz ELSE NULL END
+    WHERE project_id = ${projectId} AND source_id = ${activation.source.projectionId}
+      AND materialization_id = ${activation.materializationId}
+  `
+}
+
+export async function assertSourceRootCoverage(
+  sql: SQLClient,
+  manifest: PgOntologySourceRow
+): Promise<void> {
+  const [invalid] = await sql`
+    SELECT 1 FROM ontology_source_roots AS roots
+    WHERE project_id = ${manifest.project_id} AND source_id = ${manifest.source_id}
+      AND materialization_id = ${manifest.materialization_id}
+      AND roots.deleted = EXISTS (
+        SELECT 1 FROM ontology_source_rows AS rows
+        WHERE rows.project_id = roots.project_id AND rows.source_id = roots.source_id
+          AND rows.materialization_id = roots.materialization_id AND rows.root_sort_key = roots.root_sort_key
+      ) LIMIT 1
+  `
+  if (invalid)
+    throw new MaterializationValidationError(
+      "Source roots must have complete assertions or an explicit deletion, never both."
+    )
+}
+
+export function activeSourceRows(
+  sql: SQLClient,
+  projectId: string,
+  refs: readonly ProjectionEntityRef[]
+) {
+  const keys = refs.flatMap((entity) => {
+    const entitySortKey = utf8SortKey(projectionEntityKey(entity))
+    const roots =
+      entity.kind === "object"
+        ? [entitySortKey]
+        : [
+            entitySortKey,
+            utf8SortKey(projectionEntityKey({ kind: "object", ref: entity.ref.source })),
+          ]
+    return roots.map((rootSortKey) => ({
+      root_sort_key: rootSortKey,
+      entity_sort_key: entitySortKey,
+    }))
+  })
+  return sql<PgOntologySourceAssertionRow[]>`
+    SELECT rows.* FROM jsonb_to_recordset(${jsonParameter(sql, keys)}) AS requested(root_sort_key TEXT, entity_sort_key TEXT)
+    JOIN ontology_source_roots AS roots ON roots.project_id = ${projectId}
+      AND roots.active AND roots.root_sort_key = requested.root_sort_key
+    JOIN ontology_source_rows AS rows ON rows.project_id = roots.project_id AND rows.source_id = roots.source_id
+      AND rows.materialization_id = roots.materialization_id AND rows.entity_sort_key = requested.entity_sort_key
+      AND rows.root_sort_key = roots.root_sort_key
+  `
+}
+
+export function replacementAssertionRows(
+  sql: SQLClient,
+  input: {
+    projectId: string
+    sourceId: string
+    materializationId: string
+    incremental: boolean
+    includePrevious: boolean
+    refs: readonly ProjectionEntityRef[]
+  }
+) {
+  const keys = input.refs.map((ref) => ({
+    entity_kind: ref.kind,
+    entity_key: JSON.parse(projectionEntityKey(ref)),
+    root_keys:
+      ref.kind === "object"
+        ? [utf8SortKey(projectionEntityKey(ref))]
+        : [
+            utf8SortKey(projectionEntityKey(ref)),
+            utf8SortKey(projectionEntityKey({ kind: "object", ref: ref.ref.source })),
+          ],
+  }))
+  const previous = !input.includePrevious
+    ? sql``
+    : sql`UNION ALL
+    SELECT rows.* FROM (
+      SELECT roots.* FROM ontology_source_roots AS roots
+    WHERE roots.project_id = ${input.projectId} AND roots.source_id = ${input.sourceId}
+      AND roots.active AND roots.root_sort_key = ANY(requested.root_keys)
+      ${
+        !input.incremental
+          ? sql``
+          : sql`AND EXISTS (SELECT 1 FROM ontology_source_roots AS changed
+        WHERE changed.project_id = roots.project_id AND changed.source_id = roots.source_id
+          AND changed.materialization_id = ${input.materializationId} AND changed.root_sort_key = roots.root_sort_key)`
+      }
+
+      OFFSET 0
+    ) AS roots
+    CROSS JOIN LATERAL (
+      SELECT rows.* FROM ontology_source_rows AS rows
+      WHERE rows.project_id = roots.project_id AND rows.source_id = roots.source_id
+        AND rows.materialization_id = roots.materialization_id AND rows.root_sort_key = roots.root_sort_key
+        AND rows.entity_kind = requested.entity_kind AND rows.entity_key = requested.entity_key
+      OFFSET 0
+    ) AS rows
+  `
+  // Keep each JSON request as an indexed identity lookup even with fresh/stale statistics.
+  // Flattening this LATERAL caused a cold 10k-root run to spend seconds on each state page.
+  return sql<PgOntologySourceAssertionRow[]>`
+    SELECT selected.* FROM jsonb_to_recordset(${jsonParameter(sql, keys)})
+      AS requested(entity_kind TEXT, entity_key JSONB, root_keys TEXT[])
+    CROSS JOIN LATERAL (
+      SELECT rows.* FROM ontology_source_rows AS rows
+      WHERE rows.project_id = ${input.projectId} AND rows.source_id = ${input.sourceId}
+        AND rows.materialization_id = ${input.materializationId}
+        AND rows.entity_kind = requested.entity_kind AND rows.entity_key = requested.entity_key
+      ${previous}
+      OFFSET 0
+    ) AS selected
+  `
+}

package/src/ontology-storage/sources.ts

@@ -49,6 +49,8 @@ import {
   sourceRecord,
   toIsoString,
 } from "./shared"
+import { cleanupSourceVersions } from "./source-cleanup"
+import { assertSourceRootCoverage, stageSourceRoots } from "./source-roots"
 
 export class PgOntologySourceStorage implements OntologySourceStorage {
   constructor(private readonly runRootOperation: PgRootOperation) {}
@@ -97,14 +99,15 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
           dataset_id, dataset_version_id, dataset_version_created_at,
           projection_revision, ownership_hash, ontology_revision,
           root_count, assertion_count, created_at, ready_at, activated_at,
-          terminal_at, last_commit_id, updated_at
+          terminal_at, last_commit_id, updated_at, base_materialization_id, base_commit_id
         ) VALUES (
           ${input.projectId}, ${input.source.projectionId}, ${input.materializationId},
           ${input.execution.projectionRunId}, ${input.projectionKind}, 'replacement',
           'staging', ${input.execution.executionToken}, ${input.datasetVersion.datasetId},
           ${input.datasetVersion.versionId}, ${input.datasetVersion.createdAt},
           ${input.projectionRevision}, ${input.ownershipHash}, ${input.ontologyRevision},
-          NULL, NULL, ${input.createdAt}, NULL, NULL, NULL, NULL, ${input.createdAt}
+          NULL, NULL, ${input.createdAt}, NULL, NULL, NULL, NULL, ${input.createdAt},
+          ${input.base?.materializationId ?? null}, ${input.base?.lastCommitId ?? null}
         )
         ON CONFLICT DO NOTHING
         RETURNING *
@@ -153,6 +156,7 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
       const rows = sourceStageRows(manifest.projection_kind, input.rows)
       const existing = await this.findStageRows(sql, input, rows)
       const { pending, unchanged } = reconcileSourceStageRows(rows, existing)
+      await stageSourceRoots(sql, manifest, input)
       await this.insertStageRows(sql, input, pending)
       return { inserted: pending.length, unchanged }
     })
@@ -193,6 +197,12 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
       if (input.readyAt < toIsoString(manifest.created_at)) {
         throw new MaterializationValidationError("Source readyAt cannot precede source createdAt.")
       }
+      // A newly loaded run is absent from the planner statistics. Without refreshing them,
+      // PostgreSQL can scan the whole run for each identity instead of using its full key.
+      // Sparse deltas avoid this table-level sampling cost.
+      if (input.assertionCount >= 1_000 || input.rootCount >= 1_000) {
+        await sql`ANALYZE ontology_source_rows, ontology_source_roots`
+      }
       await this.assertReadyState(sql, manifest, input.rootCount, input.assertionCount)
       const rows = await sql<PgOntologySourceRow[]>`
         UPDATE ontology_sources
@@ -247,59 +257,7 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
       assertNonblank(input.projectId, "Terminal source cleanup project id")
       assertTimestamp(input.terminalBefore, "Terminal source cleanup cutoff", true)
       assertPositiveInteger(input.limit, "Terminal source cleanup limit")
-      const manifests = await sql<PgOntologySourceRow[]>`
-        SELECT * FROM ontology_sources
-        WHERE project_id = ${input.projectId}
-          AND status IN ('superseded', 'abandoned')
-          AND terminal_at < ${input.terminalBefore}
-        ORDER BY terminal_at, source_id, materialization_id
-        LIMIT ${input.limit}
-        FOR UPDATE SKIP LOCKED
-      `
-      let remaining = input.limit
-      let rowsDeleted = 0
-      let materializationsDeleted = 0
-      for (const manifest of manifests) {
-        if (remaining === 0) break
-        const deletedRows = await sql<{ readonly entity_sort_key: string }[]>`
-          WITH selected AS (
-            SELECT rows.ctid
-            FROM ontology_source_rows AS rows
-            WHERE rows.project_id = ${manifest.project_id}
-              AND rows.source_id = ${manifest.source_id}
-              AND rows.materialization_id = ${manifest.materialization_id}
-            ORDER BY rows.entity_sort_key
-            LIMIT ${remaining}
-            FOR UPDATE SKIP LOCKED
-          )
-          DELETE FROM ontology_source_rows AS rows
-          USING selected
-          WHERE rows.ctid = selected.ctid
-          RETURNING rows.entity_sort_key
-        `
-        rowsDeleted += deletedRows.length
-        remaining -= deletedRows.length
-        if (remaining === 0) break
-        const removed = await sql<{ readonly materialization_id: string }[]>`
-          DELETE FROM ontology_sources AS sources
-          WHERE sources.project_id = ${manifest.project_id}
-            AND sources.source_id = ${manifest.source_id}
-            AND sources.materialization_id = ${manifest.materialization_id}
-            AND sources.status IN ('superseded', 'abandoned')
-            AND sources.terminal_at = ${manifest.terminal_at}
-            AND sources.terminal_at < ${input.terminalBefore}
-            AND NOT EXISTS (
-              SELECT 1 FROM ontology_source_rows AS rows
-              WHERE rows.project_id = sources.project_id
-                AND rows.source_id = sources.source_id
-                AND rows.materialization_id = sources.materialization_id
-            )
-          RETURNING sources.materialization_id
-        `
-        materializationsDeleted += removed.length
-        remaining -= removed.length
-      }
-      return { rowsDeleted, materializationsDeleted }
+      return cleanupSourceVersions(sql, input)
     })
   }
 
@@ -408,6 +366,9 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
         `Source materialization '${manifest.materialization_id}' cannot transition to 'abandoned'.`
       )
     }
+    await sql`UPDATE ontology_source_roots SET retired_at = ${abandonedAt}
+      WHERE project_id = ${manifest.project_id} AND source_id = ${manifest.source_id}
+        AND materialization_id = ${manifest.materialization_id} AND NOT active`
     return sourceRecord(rows[0])
   }
 
@@ -426,10 +387,12 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
         readonly max_ordinal: number | string | null
       }[]
     >`
-      SELECT COUNT(*) AS assertions, COUNT(DISTINCT root_key) AS roots,
+      SELECT (SELECT COUNT(*) FROM ontology_source_rows
+        WHERE project_id = ${manifest.project_id} AND source_id = ${manifest.source_id}
+          AND materialization_id = ${manifest.materialization_id}) AS assertions, COUNT(*) AS roots,
         COUNT(DISTINCT staging_ordinal) AS ordinals,
         MIN(staging_ordinal) AS min_ordinal, MAX(staging_ordinal) AS max_ordinal
-      FROM ontology_source_rows
+      FROM ontology_source_roots
       WHERE project_id = ${manifest.project_id}
         AND source_id = ${manifest.source_id}
         AND materialization_id = ${manifest.materialization_id}
@@ -449,6 +412,7 @@ export class PgOntologySourceStorage implements OntologySourceStorage {
         "Source ready counts do not match the staged roots and assertions."
       )
     }
+    await assertSourceRootCoverage(sql, manifest)
 
     const [invalid] =
       manifest.projection_kind === "link"

package/src/pg-projection-run-storage.ts

@@ -102,7 +102,7 @@ export class PgProjectionRunStorage implements ProjectionRunStorage {
             input_exhausted = ${input.fixedBatchSize === undefined ? null : false},
             missing_target_object_type_id = ${null}, missing_target_object_id = ${null},
             missing_target_batch_ordinal = ${null}, missing_target_first_seen_at = ${null},
-            source_rows_read = ${0}, source_rows_skipped = ${0}, error = ${null}
+            source_rows_read = ${0}, source_rows_skipped = ${0}, source_changes_read = ${null}, error = ${null}
           WHERE project_id = ${input.projectId} AND id = ${input.id}
             AND status = ${"failed"} AND error->>'code' = ${"queue.enqueue_failed"}
           RETURNING *
@@ -216,7 +216,8 @@ export class PgProjectionRunStorage implements ProjectionRunStorage {
         UPDATE projection_runs
         SET
           source_rows_read = ${progress.sourceRowsRead},
-          source_rows_skipped = ${progress.sourceRowsSkipped}
+          source_rows_skipped = ${progress.sourceRowsSkipped},
+          source_changes_read = ${progress.sourceChangesRead ?? null}
         WHERE project_id = ${input.projectId}
           AND id = ${input.id}
           AND status = ${"running"}
@@ -241,6 +242,7 @@ export class PgProjectionRunStorage implements ProjectionRunStorage {
           execution_token = ${null},
           source_rows_read = ${plan.progress.sourceRowsRead},
           source_rows_skipped = ${plan.progress.sourceRowsSkipped},
+          source_changes_read = ${plan.progress.sourceChangesRead ?? null},
           input_exhausted = ${plan.inputExhausted ?? existingRow.input_exhausted},
           error = ${plan.error === undefined ? null : serializeSixbFailure(plan.error, PROJECTION_RUN_FAILURE_CODES)}::text::jsonb
         WHERE project_id = ${input.projectId}
@@ -529,6 +531,9 @@ function restoreProjectionRunRow(row: DatabaseRow): StoredProjectionRunRecord {
     progress: {
       sourceRowsRead: databaseSafeInteger(row.source_rows_read, "sourceRowsRead"),
       sourceRowsSkipped: databaseSafeInteger(row.source_rows_skipped, "sourceRowsSkipped"),
+      ...(row.source_changes_read == null
+        ? {}
+        : { sourceChangesRead: databaseSafeInteger(row.source_changes_read, "sourceChangesRead") }),
     },
     error: row.error ?? undefined,
   }
@@ -584,6 +589,7 @@ interface DatabaseRow {
   missing_target_first_seen_at: Date | string | null
   source_rows_read: number | string
   source_rows_skipped: number | string
+  source_changes_read: number | string | null
   error: JsonValue | null
 }