@sixb/pg 0.1.16 → 0.1.17

package/src/migrations/054-slim-ontology-outbox.sql

@@ -0,0 +1,108 @@
+-- An outbox event stores what changed and nothing it can be rebuilt from: delivery rebuilds the
+-- event from its commit and the execution that commit ran under. The commit also stops copying
+-- `requested_by` and `executor` from that execution (047): the execution is canonical, immutable,
+-- and already referenced by the commit.
+
+-- Every stored event must rebuild exactly as it was stored: consumers would otherwise receive a
+-- different event than the one committed. Fail rather than convert one that would not. Who asked
+-- and what wrote are derived from the execution exactly as 047 derived the commit's copies.
+DO $$
+DECLARE
+  mismatched BIGINT;
+BEGIN
+  SELECT COUNT(*) INTO mismatched
+  FROM ontology_outbox AS outbox
+  LEFT JOIN ontology_commits AS commits
+    ON commits.project_id = outbox.project_id AND commits.id = outbox.commit_id
+  LEFT JOIN executions
+    ON executions.project_id = commits.project_id AND executions.id = commits.execution_id
+  CROSS JOIN LATERAL (
+    SELECT outbox.envelope->>'type' AS type, outbox.envelope->'payload' AS payload
+  ) AS event
+  CROSS JOIN LATERAL (
+    SELECT
+      CASE
+        WHEN executions.requested_by_user_id IS NOT NULL
+          THEN jsonb_build_object('type', 'user', 'id', executions.requested_by_user_id)
+        WHEN executions.requested_by_service_account_id IS NOT NULL
+          THEN jsonb_build_object(
+            'type', 'serviceAccount', 'id', executions.requested_by_service_account_id
+          )
+      END AS requested_by,
+      CASE executions.executor_kind
+        WHEN 'request'
+          THEN jsonb_build_object('type', 'request', 'requestId', executions.executor_id)
+        WHEN 'agent' THEN jsonb_build_object('type', 'agent', 'runId', executions.executor_id)
+        WHEN 'kernel' THEN jsonb_build_object(
+          'type', 'kernel',
+          'operation',
+          CASE executions.authority_kernel_operation
+            WHEN 'ontology.recover'
+              THEN jsonb_build_object('type', 'ontology.recover', 'recoveryId', executions.executor_id)
+            ELSE jsonb_build_object(
+              'type', 'ontology.indexVectors', 'indexingId', executions.executor_id
+            )
+          END
+        )
+        ELSE jsonb_build_object(
+          'type', 'primitive',
+          'kind', executions.executor_kind,
+          'id', executions.authority_primitive_id,
+          'runId', executions.executor_id
+        )
+      END AS executor
+  ) AS attribution
+  WHERE executions.id IS NULL
+    OR outbox.envelope->>'id' IS DISTINCT FROM outbox.id
+    OR (outbox.envelope->>'commitOrdinal')::bigint IS DISTINCT FROM outbox.commit_ordinal
+    OR outbox.envelope->>'commitId' IS DISTINCT FROM commits.id
+    OR outbox.envelope->>'projectId' IS DISTINCT FROM commits.project_id
+    OR outbox.envelope->'schemaVersion' IS DISTINCT FROM '1'::jsonb
+    OR (outbox.envelope->>'occurredAt')::timestamptz IS DISTINCT FROM commits.committed_at
+    OR outbox.envelope->'origin' IS DISTINCT FROM commits.origin
+    OR outbox.envelope->>'correlationId' IS DISTINCT FROM executions.correlation_id
+    OR outbox.envelope->'requestedBy' IS DISTINCT FROM attribution.requested_by
+    OR outbox.envelope->'executor' IS DISTINCT FROM attribution.executor
+    OR outbox.envelope->>'topic' IS DISTINCT FROM CASE split_part(event.type, '.', 1)
+      WHEN 'object' THEN 'objects'
+      WHEN 'link' THEN 'links'
+      WHEN 'telemetry' THEN 'telemetry'
+    END
+    OR outbox.envelope->>'partitionKey' IS DISTINCT FROM CASE split_part(event.type, '.', 1)
+      WHEN 'object' THEN concat_ws(':', event.payload->>'objectTypeId', event.payload->>'primaryId')
+      WHEN 'link' THEN concat_ws(
+        ':', event.payload->>'sourceTypeId', event.payload->>'sourceId', event.payload->>'linkId'
+      )
+      WHEN 'telemetry' THEN concat_ws(
+        ':', event.payload->>'objectTypeId', event.payload->>'objectId', event.payload->>'propertyId'
+      )
+    END
+    OR (
+      event.type IN ('object.created', 'link.created')
+      AND event.payload->'propertyChanges' IS DISTINCT FROM COALESCE(
+        (
+          SELECT jsonb_object_agg(
+            property.key, jsonb_build_object('operation', 'created', 'after', property.value)
+          )
+          FROM jsonb_each(COALESCE(event.payload->'properties', '{}'::jsonb)) AS property
+        ),
+        '{}'::jsonb
+      )
+    );
+  IF mismatched > 0 THEN
+    RAISE EXCEPTION '[SixbPg] % outbox events would not rebuild from their commit.', mismatched;
+  END IF;
+END $$;
+
+ALTER TABLE ontology_outbox ADD COLUMN event JSONB;
+UPDATE ontology_outbox SET event = jsonb_build_object(
+  'type', envelope->'type',
+  'payload', CASE
+    WHEN envelope->>'type' IN ('object.created', 'link.created')
+      THEN (envelope->'payload') - 'propertyChanges'
+    ELSE envelope->'payload'
+  END
+);
+ALTER TABLE ontology_outbox ALTER COLUMN event SET NOT NULL, DROP COLUMN envelope;
+
+ALTER TABLE ontology_commits DROP COLUMN requested_by, DROP COLUMN executor;

package/src/migrations/055-object-query-preparation.sql

@@ -0,0 +1,11 @@
+-- Object predicates use typed property expressions, not containment on the whole JSON document.
+-- The whole-document index accelerates none of the compiler paths and amplifies every write.
+DROP INDEX IF EXISTS idx_objects_properties;
+
+-- Completion is recorded only after all non-transactional index builds and generated-column backfills finish.
+CREATE TABLE object_query_preparation (
+  singleton BOOLEAN PRIMARY KEY DEFAULT true CHECK (singleton),
+  fingerprint TEXT NOT NULL,
+  result JSONB NOT NULL,
+  prepared_at TIMESTAMPTZ NOT NULL DEFAULT now()
+);

package/src/migrations/056-commit-touches.sql

@@ -0,0 +1,41 @@
+-- Every commit records the entities whose plan inputs it changed, so a replacement plan's refresh
+-- rechecks only what the commits since it was last fresh touched, instead of the whole plan.
+-- UNLOGGED like the plans they serve: a crash truncates both together.
+CREATE UNLOGGED TABLE ontology_commit_touches (
+  project_id TEXT NOT NULL,
+  -- The transaction of the commit, compared with a plan's watermark snapshot like `xact_id` on
+  -- `ontology_commits`.
+  xact_id xid8 NOT NULL DEFAULT pg_current_xact_id(),
+  entity_kind TEXT NOT NULL CHECK (entity_kind IN ('object', 'link', 'scope')),
+  -- `objectRefKey`, `linkRefKey` or `linkScopeKey`.
+  identity_key TEXT COLLATE "C" NOT NULL
+);
+CREATE INDEX idx_ontology_commit_touches_xact ON ontology_commit_touches (project_id, xact_id);
+
+-- Touches of the project's commits before this transaction may have been purged: a plan whose
+-- watermark does not see them all checks every identity.
+CREATE UNLOGGED TABLE ontology_commit_touch_horizons (
+  project_id TEXT PRIMARY KEY,
+  xact_id xid8 NOT NULL
+);
+
+-- Existing plans predate the touches of the commits since their watermark: they plan again.
+TRUNCATE ontology_replacement_plan_work, ontology_replacement_plan_identities,
+  ontology_replacement_plans;
+
+-- Link expansion starts only from what changed since it last ran: `expansions` counts the runs,
+-- and an identity's `expand_at` is the run that has yet to start from it (an object once planned,
+-- a link once it needs a diff). `source_key` and `target_key` find a link from either endpoint.
+ALTER TABLE ontology_replacement_plans ADD COLUMN expansions INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE ontology_replacement_plan_identities
+  ADD COLUMN expand_at INTEGER,
+  ADD COLUMN source_key TEXT COLLATE "C",
+  ADD COLUMN target_key TEXT COLLATE "C";
+CREATE INDEX idx_ontology_replacement_plan_identities_expand
+  ON ontology_replacement_plan_identities (version_id, entity_kind, expand_at);
+CREATE INDEX idx_ontology_replacement_plan_identities_source
+  ON ontology_replacement_plan_identities (version_id, source_key)
+  WHERE source_key IS NOT NULL;
+CREATE INDEX idx_ontology_replacement_plan_identities_target
+  ON ontology_replacement_plan_identities (version_id, target_key)
+  WHERE target_key IS NOT NULL;

package/src/migrations/057-audio-transcription.sql

@@ -0,0 +1,10 @@
+ALTER TABLE ai_model_call_usage
+  ADD COLUMN audio_duration_ms BIGINT
+  CHECK (audio_duration_ms IS NULL OR audio_duration_ms >= 0);
+
+ALTER TABLE ai_model_call_usage
+  ADD COLUMN model_kind TEXT
+  CHECK (
+    model_kind IS NULL
+    OR model_kind IN ('language', 'image', 'video', 'embedding', 'decision', 'transcription')
+  );

package/src/migrations/058-reranking-model-kind.sql

@@ -0,0 +1,11 @@
+ALTER TABLE ai_model_call_usage
+  DROP CONSTRAINT ai_model_call_usage_model_kind_check;
+
+ALTER TABLE ai_model_call_usage
+  ADD CONSTRAINT ai_model_call_usage_model_kind_check
+  CHECK (
+    model_kind IS NULL
+    OR model_kind IN (
+      'language', 'image', 'video', 'embedding', 'decision', 'transcription', 'reranking'
+    )
+  );

package/src/migrations/059-file-download-grants.sql

@@ -0,0 +1,16 @@
+-- URLs issued by blobs.createDownloadUrl(). Only the SHA-256 of each URL's token is stored.
+CREATE TABLE file_download_grants (
+  project_id TEXT NOT NULL CHECK (length(trim(project_id)) > 0),
+  id TEXT NOT NULL CHECK (length(trim(id)) > 0),
+  token_hash TEXT NOT NULL CHECK (token_hash ~ '^[0-9a-f]{64}$'),
+  file JSONB NOT NULL CHECK (jsonb_typeof(file) = 'object'),
+  execution_id TEXT NOT NULL CHECK (length(trim(execution_id)) > 0),
+  created_at TIMESTAMPTZ NOT NULL,
+  expires_at TIMESTAMPTZ NOT NULL,
+  revoked_at TIMESTAMPTZ,
+  PRIMARY KEY (project_id, id),
+  UNIQUE (project_id, token_hash),
+  CHECK (expires_at > created_at)
+);
+
+CREATE INDEX idx_file_download_grants_expiry ON file_download_grants (project_id, expires_at);

package/src/migrations/060-drained-source-versions.sql

@@ -0,0 +1,11 @@
+-- A superseded version's `root_count` counts the roots it still holds: cleanup lowers it as it
+-- deletes them. A version that reaches zero is found through this index, instead of every
+-- superseded version being walked and probed for roots on every maintenance pass.
+UPDATE ontology_sources AS versions
+SET root_count = (SELECT count(*) FROM ontology_source_roots AS roots
+  WHERE roots.version_id = versions.version_id)
+WHERE status = 'superseded' AND root_count <> (SELECT count(*) FROM ontology_source_roots AS roots
+  WHERE roots.version_id = versions.version_id);
+CREATE INDEX idx_ontology_sources_drained
+  ON ontology_sources (project_id, terminal_at, source_id, materialization_id)
+  WHERE status = 'superseded' AND root_count = 0;

package/src/migrations.ts

@@ -115,6 +115,14 @@ import projectionRunSupersessionSql from "./migrations/051-projection-run-supers
 import compactSourceStorageSql from "./migrations/052-compact-source-storage.sql" with {
   type: "text",
 }
+import replacementPlansSql from "./migrations/053-replacement-plans.sql" with { type: "text" }
+import slimOntologyOutboxSql from "./migrations/054-slim-ontology-outbox.sql" with { type: "text" }
+import objectQueryPreparationSql from "./migrations/055-object-query-preparation.sql" with {
+  type: "text",
+}
+import commitTouchesSql from "./migrations/056-commit-touches.sql" with { type: "text" }
+import audioTranscriptionSql from "./migrations/057-audio-transcription.sql" with { type: "text" }
+import rerankingModelKindSql from "./migrations/058-reranking-model-kind.sql" with { type: "text" }
 import type { ReservedSQL, SQL, SQLClient } from "./pg-client"
 import { runPgTransactionOn, undoOnFailure, withReservedPgConnection } from "./transactions"
 
@@ -415,6 +423,14 @@ export const postgresStorageMigrations = defineMigrations<PostgresMigrationConte
     pgSql("050-directory-group-memberships", directoryGroupMembershipsSql),
     pgSql("051-projection-run-supersession", projectionRunSupersessionSql),
     pgSql("052-compact-source-storage", compactSourceStorageSql),
+    pgSql("053-replacement-plans", replacementPlansSql),
+    pgSql("054-slim-ontology-outbox", slimOntologyOutboxSql),
+    pgSql("055-object-query-preparation", objectQueryPreparationSql),
+    pgSql("056-commit-touches", commitTouchesSql),
+    pgSql("057-audio-transcription", audioTranscriptionSql),
+    pgSql("058-reranking-model-kind", rerankingModelKindSql),
+    pgSql("059-file-download-grants", fileDownloadGrantsSql),
+    pgSql("060-drained-source-versions", drainedSourceVersionsSql),
   ],
 })
 
@@ -439,3 +455,7 @@ import connectorOptionalPkceSql from "./migrations/040-connector-optional-pkce.s
   type: "text",
 }
 import fileUploadSessionsSql from "./migrations/048-file-upload-sessions.sql" with { type: "text" }
+import fileDownloadGrantsSql from "./migrations/059-file-download-grants.sql" with { type: "text" }
+import drainedSourceVersionsSql from "./migrations/060-drained-source-versions.sql" with {
+  type: "text",
+}

package/src/objects/aggregate-admission.ts

@@ -0,0 +1,21 @@
+/** Limit expensive root aggregates before they reserve a database connection. Pages can
+ * use the remaining pool capacity. No results or object revisions are retained here. */
+export class PgAggregateAdmission {
+  private active = 0
+  private readonly waiting: (() => void)[] = []
+
+  constructor(private readonly concurrency = 1) {}
+
+  async run<T>(compute: () => Promise<T>): Promise<T> {
+    if (this.active >= this.concurrency)
+      await new Promise<void>((resolve) => this.waiting.push(resolve))
+    else this.active++
+    try {
+      return await compute()
+    } finally {
+      const next = this.waiting.shift()
+      if (next) next()
+      else this.active--
+    }
+  }
+}

package/src/objects/preparation-lock.ts

@@ -0,0 +1,15 @@
+import type { ReservedSQL } from "../pg-client"
+
+/** CREATE INDEX CONCURRENTLY on expressions/partial indexes waits for older snapshots.
+ * A blocking pg_advisory_lock SELECT can hold one while waiting for the index builder,
+ * producing a cycle. Try-lock in short statements and wait outside PostgreSQL instead. */
+export async function lockPgPreparation(sql: ReservedSQL, key: string): Promise<void> {
+  let delay = 25
+  while (true) {
+    const [row] = await sql<{ acquired: boolean }[]>`
+      SELECT pg_try_advisory_lock(hashtextextended(${key}, 0)) AS acquired`
+    if (row?.acquired) return
+    await Bun.sleep(delay)
+    delay = Math.min(delay * 2, 500)
+  }
+}

package/src/objects/prepare-queries.ts

@@ -0,0 +1,181 @@
+import { createHash } from "node:crypto"
+import type { ObjectQueryIndexDefinition } from "@sixb/core"
+import type { ObjectQueryPreparationPlan, ObjectQueryPreparationResult } from "@sixb/core/storage"
+import type { ReservedSQL, SQL, SQLClient } from "../pg-client"
+import { undoOnFailure, withReservedPgConnection } from "../transactions"
+import { lockPgPreparation } from "./preparation-lock"
+import { ensurePgObjectQueryIndexes, type PgObjectQueryIndex } from "./query-indexes"
+import { type PgObjectTextCountIndex, preparePgObjectTextCounts } from "./text-count-index"
+
+/** Reconcile ontology indexes and text columns once per plan, coordinating concurrent startups. */
+export async function preparePgObjectQueries(
+  sql: SQL,
+  schema: string,
+  plan: ObjectQueryPreparationPlan
+): Promise<ObjectQueryPreparationResult> {
+  const fingerprint = preparationFingerprint(plan)
+  const current = await readCompletedPreparation(sql, fingerprint)
+  if (current) return current
+
+  // Hold one session through the entire preparation, including CONCURRENTLY builds (which
+  // cannot run in a transaction). Helpers reuse it instead of taking another pool connection.
+  return withReservedPgConnection(sql, async (connection) => {
+    const key = `${schema}.objects.query-indexes`
+    await lockPgPreparation(connection, key)
+    const unlock = () => connection`SELECT pg_advisory_unlock(hashtextextended(${key}, 0))`
+    const result = await undoOnFailure(async () => {
+      const existing = await readCompletedPreparation(connection, fingerprint)
+      if (existing) return existing
+
+      console.info(
+        `[SixbPg] Preparing query indexes for ${plan.projectId}; initial backfills may block writes.`
+      )
+      const result = await prepareStructures(connection, schema, plan)
+
+      // Record completion only after every index and generated-column backfill succeeds.
+      await connection`
+        INSERT INTO object_query_preparation (singleton, fingerprint, result)
+        VALUES (true, ${fingerprint}, ${connection.json({ ...result, warnings: [...result.warnings] })})
+        ON CONFLICT (singleton) DO UPDATE
+        SET fingerprint = excluded.fingerprint, result = excluded.result, prepared_at = now()
+      `
+      return result
+    }, unlock)
+    await unlock()
+
+    return result
+  })
+}
+
+function preparationFingerprint(plan: ObjectQueryPreparationPlan): string {
+  const objectTypes = plan.objectTypes
+    .map((type) => ({
+      objectTypeId: type.objectTypeId,
+      properties: [...type.properties].sort((a, b) => a.propertyId.localeCompare(b.propertyId)),
+      indexes: type.indexes
+        .map(normalizeIndex)
+        .sort((a, b) => JSON.stringify(a).localeCompare(JSON.stringify(b))),
+    }))
+    .sort((a, b) => a.objectTypeId.localeCompare(b.objectTypeId))
+
+  // Bump the version when physical generation changes, even if the ontology stays the same.
+  const definition = { version: 1, projectId: plan.projectId, objectTypes }
+  return createHash("sha256").update(JSON.stringify(definition)).digest("hex")
+}
+
+function normalizeIndex(index: ObjectQueryIndexDefinition) {
+  if (index.kind === "sort") {
+    return {
+      kind: index.kind,
+      fields: index.fields.map((field) => ({
+        propertyId: field.propertyId,
+        direction: field.direction ?? "asc",
+      })),
+      filters: index.filters ?? [],
+    }
+  }
+
+  return {
+    kind: index.kind,
+    propertyId: index.propertyId,
+    filters: [...(index.filters ?? [])].sort(),
+  }
+}
+
+async function readCompletedPreparation(
+  sql: SQLClient,
+  fingerprint: string
+): Promise<ObjectQueryPreparationResult | undefined> {
+  const [row] = await sql<{ result: ObjectQueryPreparationResult }[]>`
+    SELECT result FROM object_query_preparation
+    WHERE singleton AND fingerprint = ${fingerprint}
+  `
+  if (!row) return undefined
+
+  return { ...row.result, status: "current" }
+}
+
+async function prepareStructures(
+  sql: ReservedSQL,
+  schema: string,
+  plan: ObjectQueryPreparationPlan
+): Promise<ObjectQueryPreparationResult> {
+  const { indexes, textCounts, warnings } = compilePreparationPlan(plan)
+
+  if (indexes.some((index) => index.kind === "text") || textCounts.length > 0) {
+    await ensureTrigrams(sql)
+  }
+
+  const names = await ensurePgObjectQueryIndexes(sql, schema, indexes)
+  const texts = await preparePgObjectTextCounts(sql, schema, textCounts)
+  return {
+    status: "prepared",
+    objectTypes: plan.objectTypes.length,
+    indexes: new Set([...names, ...texts]).size,
+    warnings,
+  }
+}
+
+function compilePreparationPlan(plan: ObjectQueryPreparationPlan) {
+  const indexes: PgObjectQueryIndex[] = []
+  const textCounts: PgObjectTextCountIndex[] = []
+  const warnings: string[] = []
+
+  for (const type of plan.objectTypes) {
+    const scope = { projectId: plan.projectId, objectTypeId: type.objectTypeId }
+    indexes.push({ ...scope, kind: "incomingLinks" })
+
+    for (const property of type.properties) {
+      if (property.text) indexes.push({ ...scope, kind: "text", propertyId: property.propertyId })
+      if (property.sortable) {
+        indexes.push({
+          ...scope,
+          kind: "sort",
+          fields: [{ propertyId: property.propertyId, scalarKind: property.scalarKind }],
+        })
+      }
+
+      if (property.filterable && property.scalarKind === "string") {
+        indexes.push({ ...scope, kind: "filter", properties: [property.propertyId] })
+      } else if (
+        property.filterable &&
+        property.scalarKind &&
+        ["integer", "double", "boolean", "date", "timestamp", "uuid"].includes(property.scalarKind)
+      ) {
+        indexes.push({ ...scope, kind: "value", propertyId: property.propertyId })
+      } else if (property.filterable) {
+        warnings.push(
+          `${type.objectTypeId}.${property.propertyId}: no dedicated equality index for ${property.scalarKind ?? "this schema"}; queries remain exact through the general path.`
+        )
+      }
+    }
+
+    for (const index of type.indexes) {
+      if (index.kind === "sort") {
+        indexes.push({
+          ...scope,
+          kind: "sort",
+          equality: index.filters,
+          fields: index.fields.map((field) => ({
+            ...field,
+            scalarKind: type.properties.find((property) => property.propertyId === field.propertyId)
+              ?.scalarKind,
+          })),
+        })
+      } else {
+        textCounts.push({ ...scope, propertyId: index.propertyId, filters: index.filters })
+      }
+    }
+  }
+  return { indexes, textCounts, warnings }
+}
+
+// Extension creation is database-wide, whereas index preparation is schema-scoped.
+async function ensureTrigrams(sql: ReservedSQL): Promise<void> {
+  await lockPgPreparation(sql, "sixb.extension.pg_trgm")
+  const unlock = () => sql`SELECT pg_advisory_unlock(hashtextextended('sixb.extension.pg_trgm', 0))`
+  await undoOnFailure(async () => {
+    await sql.unsafe("CREATE EXTENSION IF NOT EXISTS pg_trgm WITH SCHEMA public")
+  }, unlock)
+  await unlock()
+}