@tungthedev/streams-server 0.2.0

package/src/postgres/rows.ts

@@ -0,0 +1,132 @@
+import type { Pool } from "pg";
+import type { ManifestGenerationRow } from "../store/index_store";
+import type { SegmentRow, StreamRow } from "../store/rows";
+import type { PgExecutor, PgStreamRow } from "./types";
+
+export class PostgresIndexSharedStore {
+  constructor(protected readonly pool: Pool, private readonly currentTimeMs: () => bigint) {}
+
+  nowMs(): bigint {
+    return this.currentTimeMs();
+  }
+
+  async getStream(stream: string): Promise<StreamRow | null> {
+    return getStreamWithExecutor(this.pool, stream);
+  }
+
+  async getSegmentByIndex(stream: string, segmentIndex: number): Promise<SegmentRow | null> {
+    return getSegmentByIndexWithExecutor(this.pool, stream, segmentIndex);
+  }
+
+  async countUploadedSegments(stream: string): Promise<number> {
+    const res = await this.pool.query<{ max_idx: number | string | null }>(
+      `SELECT MAX(segment_index) AS max_idx FROM segments WHERE stream = $1 AND r2_etag IS NOT NULL;`,
+      [stream]
+    );
+    const maxIdx = res.rows[0]?.max_idx == null ? -1 : Number(res.rows[0]!.max_idx);
+    return maxIdx >= 0 ? maxIdx + 1 : 0;
+  }
+
+  async countSegmentsForStream(stream: string): Promise<number> {
+    const res = await this.pool.query<{ count: string }>(`SELECT COUNT(*) AS count FROM segments WHERE stream = $1;`, [stream]);
+    return Number(res.rows[0]?.count ?? 0);
+  }
+
+  async getManifestRow(stream: string): Promise<ManifestGenerationRow> {
+    return getManifestGenerationWithExecutor(this.pool, stream);
+  }
+}
+
+export async function getStreamWithExecutor(executor: PgExecutor, stream: string): Promise<StreamRow | null> {
+  const res = await executor.query<PgStreamRow>(`SELECT * FROM streams WHERE stream = $1;`, [stream]);
+  return res.rows[0] ? coerceStreamRow(res.rows[0]) : null;
+}
+
+export async function getSegmentByIndexWithExecutor(executor: PgExecutor, stream: string, segmentIndex: number): Promise<SegmentRow | null> {
+  const res = await executor.query(
+    `SELECT segment_id, stream, segment_index, start_offset, end_offset, block_count, last_append_ms,
+            payload_bytes, size_bytes, local_path, created_at_ms, uploaded_at_ms, r2_etag
+     FROM segments
+     WHERE stream = $1 AND segment_index = $2
+     LIMIT 1;`,
+    [stream, segmentIndex]
+  );
+  return res.rows[0] ? coerceSegmentRow(res.rows[0]) : null;
+}
+
+async function getManifestGenerationWithExecutor(executor: PgExecutor, stream: string): Promise<ManifestGenerationRow> {
+  const res = await executor.query<{ generation: number | string }>(`SELECT generation FROM manifests WHERE stream = $1 LIMIT 1;`, [stream]);
+  if (res.rows[0]) return { generation: Number(res.rows[0].generation) };
+  await executor.query(
+    `INSERT INTO manifests(stream, generation, uploaded_generation, last_uploaded_at_ms, last_uploaded_etag, last_uploaded_size_bytes)
+     VALUES($1, 0, 0, NULL, NULL, NULL)
+     ON CONFLICT(stream) DO NOTHING;`,
+    [stream]
+  );
+  return { generation: 0 };
+}
+
+export function coerceSegmentRow(row: any): SegmentRow {
+  return {
+    segment_id: String(row.segment_id),
+    stream: String(row.stream),
+    segment_index: Number(row.segment_index),
+    start_offset: toBigInt(row.start_offset),
+    end_offset: toBigInt(row.end_offset),
+    block_count: Number(row.block_count),
+    last_append_ms: toBigInt(row.last_append_ms),
+    payload_bytes: toBigInt(row.payload_bytes),
+    size_bytes: Number(row.size_bytes),
+    local_path: String(row.local_path),
+    created_at_ms: toBigInt(row.created_at_ms),
+    uploaded_at_ms: row.uploaded_at_ms == null ? null : toBigInt(row.uploaded_at_ms),
+    r2_etag: row.r2_etag == null ? null : String(row.r2_etag),
+  };
+}
+
+function coerceStreamRow(row: PgStreamRow): StreamRow {
+  const walRows = toBigInt(row.wal_rows);
+  const walBytes = toBigInt(row.wal_bytes);
+  const lastAppendMs = toBigInt(row.last_append_ms);
+  return {
+    stream: String(row.stream),
+    created_at_ms: toBigInt(row.created_at_ms),
+    updated_at_ms: toBigInt(row.updated_at_ms),
+    content_type: String(row.content_type),
+    profile: row.profile == null ? null : String(row.profile),
+    stream_seq: row.stream_seq == null ? null : String(row.stream_seq),
+    closed: Number(row.closed),
+    closed_producer_id: row.closed_producer_id == null ? null : String(row.closed_producer_id),
+    closed_producer_epoch: row.closed_producer_epoch == null ? null : Number(row.closed_producer_epoch),
+    closed_producer_seq: row.closed_producer_seq == null ? null : Number(row.closed_producer_seq),
+    ttl_seconds: row.ttl_seconds == null ? null : Number(row.ttl_seconds),
+    epoch: Number(row.epoch),
+    next_offset: toBigInt(row.next_offset),
+    sealed_through: row.sealed_through == null ? -1n : toBigInt(row.sealed_through),
+    uploaded_through: row.uploaded_through == null ? -1n : toBigInt(row.uploaded_through),
+    uploaded_segment_count: row.uploaded_segment_count == null ? 0 : Number(row.uploaded_segment_count),
+    pending_rows: row.pending_rows == null ? walRows : toBigInt(row.pending_rows),
+    pending_bytes: row.pending_bytes == null ? walBytes : toBigInt(row.pending_bytes),
+    logical_size_bytes: toBigInt(row.logical_size_bytes),
+    wal_rows: walRows,
+    wal_bytes: walBytes,
+    last_append_ms: lastAppendMs,
+    last_segment_cut_ms: row.last_segment_cut_ms == null ? lastAppendMs : toBigInt(row.last_segment_cut_ms),
+    segment_in_progress: row.segment_in_progress == null ? 0 : Number(row.segment_in_progress),
+    expires_at_ms: row.expires_at_ms == null ? null : toBigInt(row.expires_at_ms),
+    stream_flags: Number(row.stream_flags),
+  };
+}
+
+export function pgInt(value: bigint): string {
+  return value.toString();
+}
+
+export function toBigInt(value: unknown): bigint {
+  return typeof value === "bigint" ? value : BigInt(String(value));
+}
+
+export function toBytes(value: unknown): Uint8Array {
+  if (value instanceof Uint8Array) return value;
+  return new Uint8Array(value as ArrayBuffer);
+}

package/src/postgres/schema.ts

@@ -0,0 +1,355 @@
+import type { Pool } from "pg";
+import { dsError } from "../util/ds_error";
+
+export const POSTGRES_SCHEMA_VERSION = 1;
+
+type PgSchemaExecutor = Pick<Pool, "query">;
+
+export type PostgresMigrationOptions = {
+  fullMode?: boolean;
+};
+
+export async function migratePostgresStore(pool: Pool, opts: PostgresMigrationOptions = {}): Promise<void> {
+  const client = await pool.connect();
+  try {
+    await client.query("BEGIN");
+    await ensureSchemaVersionTable(client);
+    const currentVersion = await readPostgresSchemaVersion(client);
+    if (currentVersion != null && currentVersion !== POSTGRES_SCHEMA_VERSION) {
+      throw dsError(`postgres schema version ${currentVersion} is not supported by version ${POSTGRES_SCHEMA_VERSION}`);
+    }
+    await installWalControlPlaneSchema(client);
+    if (opts.fullMode) {
+      await installFullModeSegmentSchema(client);
+      await installFullModeIndexSchema(client);
+      await installFullModeTouchSchema(client);
+      await installFullModeStatsSchema(client);
+    }
+    await setPostgresSchemaVersion(client, POSTGRES_SCHEMA_VERSION);
+    await client.query("COMMIT");
+  } catch (error) {
+    await client.query("ROLLBACK").catch(() => {});
+    throw error;
+  } finally {
+    client.release();
+  }
+}
+
+export async function readPostgresSchemaVersion(executor: PgSchemaExecutor): Promise<number | null> {
+  const table = await executor.query<{ exists: string | null }>(`SELECT to_regclass('schema_version')::text AS exists;`);
+  if (!table.rows[0]?.exists) return null;
+  const rows = await executor.query<{ version: number | string }>(`SELECT version FROM schema_version;`);
+  if (rows.rows.length === 0) return null;
+  if (rows.rows.length > 1) throw dsError("postgres schema_version must contain exactly one row");
+  return Number(rows.rows[0]!.version);
+}
+
+async function ensureSchemaVersionTable(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS schema_version (
+      version integer NOT NULL
+    );
+  `);
+}
+
+async function setPostgresSchemaVersion(executor: PgSchemaExecutor, version: number): Promise<void> {
+  await executor.query(`DELETE FROM schema_version;`);
+  await executor.query(`INSERT INTO schema_version(version) VALUES($1);`, [version]);
+}
+
+async function installWalControlPlaneSchema(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`
+      CREATE TABLE IF NOT EXISTS streams (
+        stream text PRIMARY KEY,
+        created_at_ms bigint NOT NULL,
+        updated_at_ms bigint NOT NULL,
+        content_type text NOT NULL,
+        profile text NULL,
+        stream_seq text NULL,
+        closed integer NOT NULL DEFAULT 0,
+        closed_producer_id text NULL,
+        closed_producer_epoch integer NULL,
+        closed_producer_seq integer NULL,
+        ttl_seconds integer NULL,
+        epoch integer NOT NULL DEFAULT 0,
+        next_offset bigint NOT NULL DEFAULT 0,
+        logical_size_bytes bigint NOT NULL DEFAULT 0,
+        wal_rows bigint NOT NULL DEFAULT 0,
+        wal_bytes bigint NOT NULL DEFAULT 0,
+        last_append_ms bigint NOT NULL,
+        expires_at_ms bigint NULL,
+        stream_flags integer NOT NULL DEFAULT 0
+      );
+    `);
+  await executor.query(`
+      CREATE TABLE IF NOT EXISTS wal (
+        stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+        "offset" bigint NOT NULL,
+        ts_ms bigint NOT NULL,
+        payload bytea NOT NULL,
+        payload_len integer NOT NULL,
+        routing_key bytea NULL,
+        content_type text NULL,
+        flags integer NOT NULL DEFAULT 0,
+        PRIMARY KEY (stream, "offset")
+      );
+    `);
+  await executor.query(`
+      CREATE TABLE IF NOT EXISTS schemas (
+        stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+        schema_json text NOT NULL,
+        updated_at_ms bigint NOT NULL
+      );
+    `);
+  await executor.query(`
+      CREATE TABLE IF NOT EXISTS stream_profiles (
+        stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+        profile_json text NOT NULL,
+        updated_at_ms bigint NOT NULL
+      );
+    `);
+  await executor.query(`
+      CREATE TABLE IF NOT EXISTS producer_state (
+        stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+        producer_id text NOT NULL,
+        epoch integer NOT NULL,
+        last_seq integer NOT NULL,
+        updated_at_ms bigint NOT NULL,
+        PRIMARY KEY (stream, producer_id)
+      );
+    `);
+  await executor.query(`CREATE INDEX IF NOT EXISTS wal_stream_offset_idx ON wal(stream, "offset");`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS wal_stream_routing_offset_idx ON wal(stream, routing_key, "offset");`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS streams_updated_at_idx ON streams(updated_at_ms);`);
+}
+
+async function installFullModeSegmentSchema(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS sealed_through bigint NOT NULL DEFAULT -1;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS uploaded_through bigint NOT NULL DEFAULT -1;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS uploaded_segment_count integer NOT NULL DEFAULT 0;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS pending_rows bigint NOT NULL DEFAULT 0;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS pending_bytes bigint NOT NULL DEFAULT 0;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS last_segment_cut_ms bigint NOT NULL DEFAULT 0;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS segment_in_progress integer NOT NULL DEFAULT 0;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS segment_claim_token text NULL;`);
+  await executor.query(`ALTER TABLE streams ADD COLUMN IF NOT EXISTS segment_claimed_at_ms bigint NULL;`);
+  await executor.query(`
+    UPDATE streams
+    SET pending_rows = wal_rows,
+        pending_bytes = wal_bytes,
+        last_segment_cut_ms = last_append_ms
+    WHERE sealed_through = -1
+      AND uploaded_through = -1
+      AND uploaded_segment_count = 0
+      AND segment_in_progress = 0
+      AND segment_claim_token IS NULL
+      AND pending_rows = 0
+      AND pending_bytes = 0;
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS segments (
+      segment_id text PRIMARY KEY,
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      segment_index integer NOT NULL,
+      start_offset bigint NOT NULL,
+      end_offset bigint NOT NULL,
+      block_count integer NOT NULL,
+      last_append_ms bigint NOT NULL,
+      payload_bytes bigint NOT NULL DEFAULT 0,
+      size_bytes integer NOT NULL,
+      local_path text NOT NULL,
+      created_at_ms bigint NOT NULL,
+      uploaded_at_ms bigint NULL,
+      r2_etag text NULL,
+      UNIQUE(stream, segment_index)
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS stream_segment_meta (
+      stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+      segment_count integer NOT NULL,
+      segment_offsets bytea NOT NULL,
+      segment_blocks bytea NOT NULL,
+      segment_last_ts bytea NOT NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS manifests (
+      stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+      generation integer NOT NULL,
+      uploaded_generation integer NOT NULL,
+      last_uploaded_at_ms bigint NULL,
+      last_uploaded_etag text NULL,
+      last_uploaded_size_bytes bigint NULL
+    );
+  `);
+  await executor.query(`ALTER TABLE schemas ADD COLUMN IF NOT EXISTS uploaded_size_bytes bigint NOT NULL DEFAULT 0;`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS streams_pending_bytes_idx ON streams(pending_bytes);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS streams_last_cut_idx ON streams(last_segment_cut_ms);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS streams_inprog_pending_idx ON streams(segment_in_progress, pending_bytes, last_segment_cut_ms);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS segments_stream_start_idx ON segments(stream, start_offset);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS segments_pending_upload_idx ON segments(uploaded_at_ms);`);
+}
+
+async function installFullModeIndexSchema(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS index_state (
+      stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+      index_secret bytea NOT NULL,
+      indexed_through integer NOT NULL,
+      updated_at_ms bigint NOT NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS index_runs (
+      run_id text PRIMARY KEY,
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      level integer NOT NULL,
+      start_segment integer NOT NULL,
+      end_segment integer NOT NULL,
+      object_key text NOT NULL,
+      size_bytes bigint NOT NULL,
+      filter_len integer NOT NULL,
+      record_count integer NOT NULL,
+      retired_gen integer NULL,
+      retired_at_ms bigint NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS secondary_index_state (
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      index_name text NOT NULL,
+      index_secret bytea NOT NULL,
+      config_hash text NOT NULL,
+      indexed_through integer NOT NULL,
+      updated_at_ms bigint NOT NULL,
+      PRIMARY KEY(stream, index_name)
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS secondary_index_runs (
+      run_id text PRIMARY KEY,
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      index_name text NOT NULL,
+      level integer NOT NULL,
+      start_segment integer NOT NULL,
+      end_segment integer NOT NULL,
+      object_key text NOT NULL,
+      size_bytes bigint NOT NULL,
+      filter_len integer NOT NULL,
+      record_count integer NOT NULL,
+      retired_gen integer NULL,
+      retired_at_ms bigint NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS lexicon_index_state (
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      source_kind text NOT NULL,
+      source_name text NOT NULL,
+      indexed_through integer NOT NULL,
+      updated_at_ms bigint NOT NULL,
+      PRIMARY KEY(stream, source_kind, source_name)
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS lexicon_index_runs (
+      run_id text PRIMARY KEY,
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      source_kind text NOT NULL,
+      source_name text NOT NULL,
+      level integer NOT NULL,
+      start_segment integer NOT NULL,
+      end_segment integer NOT NULL,
+      object_key text NOT NULL,
+      size_bytes bigint NOT NULL,
+      record_count integer NOT NULL,
+      retired_gen integer NULL,
+      retired_at_ms bigint NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS search_companion_plans (
+      stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+      generation integer NOT NULL,
+      plan_hash text NOT NULL,
+      plan_json text NOT NULL,
+      updated_at_ms bigint NOT NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS search_segment_companions (
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      segment_index integer NOT NULL,
+      object_key text NOT NULL,
+      plan_generation integer NOT NULL,
+      sections_json text NOT NULL,
+      section_sizes_json text NOT NULL,
+      size_bytes bigint NOT NULL,
+      primary_timestamp_min_ms bigint NULL,
+      primary_timestamp_max_ms bigint NULL,
+      updated_at_ms bigint NOT NULL,
+      PRIMARY KEY(stream, segment_index)
+    );
+  `);
+  await executor.query(`CREATE INDEX IF NOT EXISTS index_runs_active_idx ON index_runs(stream, retired_gen, level, start_segment);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS index_runs_retired_idx ON index_runs(stream, retired_gen, retired_at_ms);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS secondary_index_runs_active_idx ON secondary_index_runs(stream, index_name, retired_gen, level, start_segment);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS secondary_index_runs_retired_idx ON secondary_index_runs(stream, index_name, retired_gen, retired_at_ms);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS lexicon_index_runs_active_idx ON lexicon_index_runs(stream, source_kind, source_name, retired_gen, level, start_segment);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS lexicon_index_runs_retired_idx ON lexicon_index_runs(stream, source_kind, source_name, retired_gen, retired_at_ms);`);
+  await executor.query(`CREATE INDEX IF NOT EXISTS search_segment_companions_plan_idx ON search_segment_companions(stream, plan_generation, segment_index);`);
+}
+
+async function installFullModeTouchSchema(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS stream_touch_state (
+      stream text PRIMARY KEY REFERENCES streams(stream) ON DELETE CASCADE,
+      processed_through bigint NOT NULL,
+      updated_at_ms bigint NOT NULL
+    );
+  `);
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS live_templates (
+      stream text NOT NULL REFERENCES streams(stream) ON DELETE CASCADE,
+      template_id text NOT NULL,
+      entity text NOT NULL,
+      fields_json text NOT NULL,
+      encodings_json text NOT NULL,
+      state text NOT NULL,
+      created_at_ms bigint NOT NULL,
+      last_seen_at_ms bigint NOT NULL,
+      inactivity_ttl_ms bigint NOT NULL,
+      active_from_source_offset bigint NOT NULL,
+      retired_at_ms bigint NULL,
+      retired_reason text NULL,
+      PRIMARY KEY(stream, template_id)
+    );
+  `);
+  await executor.query(`
+    CREATE INDEX IF NOT EXISTS live_templates_stream_entity_state_last_seen_idx
+      ON live_templates(stream, entity, state, last_seen_at_ms);
+  `);
+  await executor.query(`
+    CREATE INDEX IF NOT EXISTS live_templates_stream_state_last_seen_idx
+      ON live_templates(stream, state, last_seen_at_ms);
+  `);
+}
+
+async function installFullModeStatsSchema(executor: PgSchemaExecutor): Promise<void> {
+  await executor.query(`
+    CREATE TABLE IF NOT EXISTS objectstore_request_counts (
+      stream_hash text NOT NULL,
+      artifact text NOT NULL,
+      op text NOT NULL,
+      count bigint NOT NULL DEFAULT 0,
+      bytes bigint NOT NULL DEFAULT 0,
+      updated_at_ms bigint NOT NULL,
+      PRIMARY KEY(stream_hash, artifact, op)
+    );
+  `);
+  await executor.query(`
+    CREATE INDEX IF NOT EXISTS objectstore_request_counts_stream_hash_idx
+      ON objectstore_request_counts(stream_hash, updated_at_ms);
+  `);
+}