@zmdb/cockroach 1.0.0-beta.1

package/package.json

@@ -0,0 +1,54 @@
+{
+  "name": "@zmdb/cockroach",
+  "version": "1.0.0-beta.1",
+  "description": "CockroachDB vertical for zmdb: PostgreSQL-family dialect overrides, migrations, catalog introspection, retries, and a pg-protocol driver.",
+  "keywords": [
+    "cockroachdb",
+    "database",
+    "migrations",
+    "orm",
+    "postgresql",
+    "typescript",
+    "zmdb"
+  ],
+  "homepage": "https://github.com/ambasta/zmdb#readme",
+  "bugs": {
+    "url": "https://github.com/ambasta/zmdb/issues"
+  },
+  "license": "GPL-3.0-or-later",
+  "author": "zmdb contributors",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/ambasta/zmdb.git",
+    "directory": "packages/cockroach"
+  },
+  "type": "module",
+  "sideEffects": false,
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "publishConfig": {
+    "access": "public",
+    "tag": "beta"
+  },
+  "scripts": {
+    "build": "node ../../scripts/build-package.mjs",
+    "test": "vitest run"
+  },
+  "dependencies": {
+    "@zmdb/migrations": "1.0.0-beta.1",
+    "@zmdb/postgres": "1.0.0-beta.1"
+  },
+  "peerDependencies": {
+    "@zmdb/orm": "1.0.0-beta.1",
+    "@zmdb/sql": "1.0.0-beta.1"
+  },
+  "engines": {
+    "node": ">=26"
+  },
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts"
+}

package/src/index.ts

@@ -0,0 +1,146 @@
+import { type DatabaseVertical, type TransactionalDriver } from '@zmdb/orm';
+import { postgres, postgresFamilyDriver, type PgConnection, type PgOptions, type PgQueryable } from '@zmdb/postgres';
+import { extendSqlDialect, UnsupportedFeatureError, type SqlDialect } from '@zmdb/sql';
+
+import { cockroachIntrospector } from './introspect.js';
+import { COCKROACH_TYPE_OVERRIDES, cockroachMigrations } from './migrations.js';
+
+export type { PgConnection, PgOptions, PgQueryable };
+export { cockroachIntrospector, cockroachMigrations };
+
+interface PgPoolLike extends PgQueryable {
+  readonly totalCount: number;
+  readonly idleCount: number;
+  connect(): Promise<PgConnection>;
+}
+
+type PgQueryConfig = {
+  readonly name?: string;
+  readonly queryMode?: 'extended';
+  readonly text: string;
+  readonly values?: readonly unknown[];
+};
+
+function isPoolLike(client: PgQueryable): client is PgPoolLike {
+  return (
+    typeof client.connect === 'function' &&
+    typeof Reflect.get(client, 'totalCount') === 'number' &&
+    typeof Reflect.get(client, 'idleCount') === 'number'
+  );
+}
+
+function normalizeBackendPid(
+  input: string | PgQueryConfig,
+  result: { rows: Record<string, unknown>[] },
+): { rows: Record<string, unknown>[] } {
+  const text = typeof input === 'string' ? input : input.text;
+  if (text.trim().toLowerCase() !== 'select pg_backend_pid() as pid') return result;
+  return {
+    rows: result.rows.map(row => {
+      const pid = Reflect.get(row, 'pid');
+      if (typeof pid !== 'string' || !/^\d+$/.test(pid)) return row;
+      const numeric = Number(pid);
+      return Number.isSafeInteger(numeric) && numeric > 0 ? { ...row, pid: numeric } : row;
+    }),
+  };
+}
+
+class CockroachClient implements PgConnection {
+  private readonly target: PgConnection;
+
+  constructor(target: PgConnection) {
+    this.target = target;
+  }
+
+  query(text: string, params?: readonly unknown[]): Promise<{ rows: Record<string, unknown>[] }>;
+  query(config: PgQueryConfig): Promise<{ rows: Record<string, unknown>[] }>;
+  async query(
+    input: string | PgQueryConfig,
+    params?: readonly unknown[],
+  ): Promise<{ rows: Record<string, unknown>[] }> {
+    const result = typeof input === 'string' ? await this.target.query(input, params) : await this.target.query(input);
+    return normalizeBackendPid(input, result);
+  }
+
+  release(): void {
+    const release = this.target.release;
+    if (release !== undefined) release.call(this.target);
+  }
+}
+
+class CockroachQueryable implements PgQueryable {
+  private readonly target: PgQueryable;
+
+  constructor(target: PgQueryable) {
+    this.target = target;
+  }
+
+  query(text: string, params?: readonly unknown[]): Promise<{ rows: Record<string, unknown>[] }>;
+  query(config: PgQueryConfig): Promise<{ rows: Record<string, unknown>[] }>;
+  async query(
+    input: string | PgQueryConfig,
+    params?: readonly unknown[],
+  ): Promise<{ rows: Record<string, unknown>[] }> {
+    const result = typeof input === 'string' ? await this.target.query(input, params) : await this.target.query(input);
+    return normalizeBackendPid(input, result);
+  }
+}
+
+class CockroachPool extends CockroachQueryable implements PgPoolLike {
+  private readonly pool: PgPoolLike;
+
+  constructor(pool: PgPoolLike) {
+    super(pool);
+    this.pool = pool;
+  }
+
+  get totalCount(): number {
+    return this.pool.totalCount;
+  }
+
+  get idleCount(): number {
+    return this.pool.idleCount;
+  }
+
+  async connect(): Promise<PgConnection> {
+    return new CockroachClient(await this.pool.connect());
+  }
+}
+
+function cockroachQueryable(client: PgQueryable): PgQueryable {
+  return isPoolLike(client) ? new CockroachPool(client) : new CockroachQueryable(client);
+}
+
+export const cockroach: SqlDialect<'cockroach'> = extendSqlDialect(postgres, {
+  name: 'cockroach',
+  traits: {
+    types: COCKROACH_TYPE_OVERRIDES,
+    fts: 'none',
+    retryableCodes: ['40001'],
+    vectorDistance: false,
+    spatialPredicates: false,
+  },
+  capabilities: {
+    transactionalDdl: false,
+    rowLevelSecurity: false,
+    cancellation: false,
+  },
+  migrations: cockroachMigrations,
+  introspector: cockroachIntrospector,
+});
+
+export function cockroachDriver(client: PgQueryable, options?: PgOptions): TransactionalDriver<'cockroach'> {
+  if (options?.cancelVia !== undefined) {
+    throw new UnsupportedFeatureError(
+      'server-side cancellation',
+      'cockroach',
+      'cockroach does not provide PostgreSQL pg_cancel_backend(); omit PgOptions.cancelVia',
+    );
+  }
+  return postgresFamilyDriver(cockroach, cockroachQueryable(client), options);
+}
+
+export const cockroachVertical: DatabaseVertical<'cockroach', PgQueryable, PgOptions> = Object.freeze({
+  dialect: cockroach,
+  driver: cockroachDriver,
+});

package/src/introspect.ts

@@ -0,0 +1,289 @@
+import {
+  CatalogRowError,
+  integerField,
+  query,
+  textField,
+  type CatalogColumnSnapshot,
+  type CatalogIndexColumn,
+  type CatalogIndexSnapshot,
+  type CatalogSchemaSnapshot,
+  type CatalogTableSnapshot,
+  type IntrospectionDriver,
+} from '@zmdb/migrations/introspect/runtime';
+import { postgresFamilyIntrospector } from '@zmdb/postgres';
+import { type IntrospectOptions } from '@zmdb/sql';
+
+interface CockroachIndexRow {
+  readonly name: string;
+  readonly nonUnique: boolean;
+  readonly sequence: number;
+  readonly column: string;
+  readonly definition: string;
+  readonly storing: boolean;
+  readonly implicit: boolean;
+}
+
+function quoteIdentifier(identifier: string): string {
+  return `"${identifier.replaceAll('"', '""')}"`;
+}
+
+function selectedSchemas(options: IntrospectOptions | undefined): readonly string[] {
+  const schemas = options?.schemas ?? ['public'];
+  return schemas.length === 0 ? ['public'] : [...new Set(schemas)].toSorted();
+}
+
+function placeholders(count: number): string {
+  return Array.from({ length: count }, (_, index) => `$${String(index + 1)}`).join(', ');
+}
+
+function booleanField(row: Readonly<Record<string, unknown>>, field: string, catalog: string, index: number): boolean {
+  const value = Reflect.get(row, field);
+  if (typeof value === 'boolean') return value;
+  if (value === 't' || value === 'true' || value === 1 || value === '1') return true;
+  if (value === 'f' || value === 'false' || value === 0 || value === '0') return false;
+  throw new CatalogRowError(catalog, index, field, 'a boolean', value);
+}
+
+function parseIndexRow(row: Readonly<Record<string, unknown>>, index: number): CockroachIndexRow {
+  const catalog = 'cockroach SHOW INDEXES';
+  return {
+    name: textField(row, 'index_name', catalog, index),
+    nonUnique: booleanField(row, 'non_unique', catalog, index),
+    sequence: integerField(row, 'seq_in_index', catalog, index),
+    column: textField(row, 'column_name', catalog, index),
+    definition: textField(row, 'definition', catalog, index),
+    storing: booleanField(row, 'storing', catalog, index),
+    implicit: booleanField(row, 'implicit', catalog, index),
+  };
+}
+
+function unquoteIdentifier(value: string): string {
+  return value.startsWith('"') && value.endsWith('"') ? value.slice(1, -1).replaceAll('""', '"') : value;
+}
+
+function balancedOuterParentheses(value: string): boolean {
+  if (!value.startsWith('(') || !value.endsWith(')')) return false;
+  let depth = 0;
+  let quoteCharacter: "'" | '"' | undefined;
+  for (let index = 0; index < value.length; index += 1) {
+    const character = value[index];
+    if (quoteCharacter !== undefined) {
+      if (character === quoteCharacter) {
+        if (value[index + 1] === quoteCharacter) index += 1;
+        else quoteCharacter = undefined;
+      }
+      continue;
+    }
+    if (character === "'" || character === '"') {
+      quoteCharacter = character;
+      continue;
+    }
+    if (character === '(') depth += 1;
+    if (character === ')') depth -= 1;
+    if (depth === 0 && index < value.length - 1) return false;
+  }
+  return depth === 0;
+}
+
+function stripOuterParentheses(value: string): string {
+  let current = value.trim();
+  while (balancedOuterParentheses(current)) current = current.slice(1, -1).trim();
+  return current;
+}
+
+function indexColumn(row: CockroachIndexRow, tableColumns: ReadonlySet<string>): CatalogIndexColumn {
+  if (!row.column.startsWith('crdb_internal_') && tableColumns.has(row.column)) return row.column;
+  const definition = row.definition.trim();
+  const identifier = /^("(?:[^"]|"")*"|[A-Za-z_][A-Za-z0-9_$]*)$/.exec(definition)?.[1];
+  if (identifier !== undefined) {
+    const column = unquoteIdentifier(identifier);
+    if (tableColumns.has(column)) return column;
+  }
+  return { expr: stripOuterParentheses(definition) };
+}
+
+function identifierAt(value: string, start: number): { readonly name: string; readonly end: number } | undefined {
+  if (value[start] === '"') {
+    let index = start + 1;
+    let name = '';
+    while (index < value.length) {
+      const character = value[index];
+      if (character === '"') {
+        if (value[index + 1] === '"') {
+          name += '"';
+          index += 2;
+          continue;
+        }
+        return { name, end: index + 1 };
+      }
+      if (character === undefined) return undefined;
+      name += character;
+      index += 1;
+    }
+    return undefined;
+  }
+  const match = /^[A-Za-z_][A-Za-z0-9_$]*/.exec(value.slice(start));
+  const name = match?.[0];
+  return name === undefined ? undefined : { name, end: start + name.length };
+}
+
+function matchingParenthesis(value: string, open: number): number | undefined {
+  let depth = 0;
+  let quoteCharacter: "'" | '"' | undefined;
+  for (let index = open; index < value.length; index += 1) {
+    const character = value[index];
+    if (quoteCharacter !== undefined) {
+      if (character === quoteCharacter) {
+        if (value[index + 1] === quoteCharacter) index += 1;
+        else quoteCharacter = undefined;
+      }
+      continue;
+    }
+    if (character === "'" || character === '"') {
+      quoteCharacter = character;
+      continue;
+    }
+    if (character === '(') depth += 1;
+    if (character === ')') {
+      depth -= 1;
+      if (depth === 0) return index;
+    }
+  }
+  return undefined;
+}
+
+function indexPredicates(createStatement: string): ReadonlyMap<string, string> {
+  const predicates = new Map<string, string>();
+  for (const rawLine of createStatement.split('\n')) {
+    const line = rawLine.trim().replace(/,$/, '');
+    const prefix = /^(?:UNIQUE\s+)?INDEX\s+/i.exec(line)?.[0];
+    if (prefix === undefined) continue;
+    const identifier = identifierAt(line, prefix.length);
+    if (identifier === undefined) continue;
+    const open = line.indexOf('(', identifier.end);
+    if (open === -1) continue;
+    const close = matchingParenthesis(line, open);
+    if (close === undefined) continue;
+    const suffix = line.slice(close + 1);
+    const where = /\sWHERE\s+(.+)$/i.exec(suffix)?.[1]?.trim();
+    if (where !== undefined && where.length > 0) predicates.set(identifier.name, where);
+  }
+  return predicates;
+}
+
+function indexes(
+  rows: readonly CockroachIndexRow[],
+  primaryIndexes: ReadonlySet<string>,
+  tableColumns: ReadonlySet<string>,
+  predicates: ReadonlyMap<string, string>,
+): readonly CatalogIndexSnapshot[] {
+  const grouped = new Map<string, CockroachIndexRow[]>();
+  for (const row of rows) {
+    if (primaryIndexes.has(row.name) || row.storing || row.implicit) continue;
+    const values = grouped.get(row.name);
+    if (values === undefined) grouped.set(row.name, [row]);
+    else values.push(row);
+  }
+  return [...grouped]
+    .map(([name, values]) => {
+      const sorted = values.toSorted((left, right) => left.sequence - right.sequence);
+      const first = sorted[0];
+      if (first === undefined) throw new Error(`cockroach index "${name}" has no key columns`);
+      const where = predicates.get(name);
+      return {
+        name,
+        columns: sorted.map(row => indexColumn(row, tableColumns)),
+        unique: !first.nonUnique,
+        ...(where === undefined ? {} : { where }),
+      };
+    })
+    .toSorted((left, right) => left.name.localeCompare(right.name));
+}
+
+async function schemasByTable(
+  driver: IntrospectionDriver,
+  options: IntrospectOptions | undefined,
+): Promise<ReadonlyMap<string, string>> {
+  const schemas = selectedSchemas(options);
+  const rows = await driver.execute(
+    query(
+      `SELECT table_schema, table_name FROM information_schema.tables ` +
+        `WHERE table_type = 'BASE TABLE' AND table_schema IN (${placeholders(schemas.length)}) ` +
+        'ORDER BY table_schema, table_name',
+      schemas,
+    ),
+  );
+  return new Map(
+    rows.map((row, index) => [
+      textField(row, 'table_name', 'cockroach information_schema.tables', index),
+      textField(row, 'table_schema', 'cockroach information_schema.tables', index),
+    ]),
+  );
+}
+
+function cockroachColumn(column: CatalogColumnSnapshot): CatalogColumnSnapshot {
+  const defaultValue = column.default?.trim().toLowerCase();
+  if (
+    column.type !== 'bigint' ||
+    defaultValue === undefined ||
+    !/^unique_rowid\(\)(?:::[a-z0-9_ ]+)?$/.test(defaultValue)
+  ) {
+    return column;
+  }
+  const { default: ignoredDefault, ...withoutDefault } = column;
+  void ignoredDefault;
+  return { ...withoutDefault, type: 'serial' };
+}
+
+async function cockroachTable(
+  driver: IntrospectionDriver,
+  schema: string,
+  table: CatalogTableSnapshot,
+): Promise<CatalogTableSnapshot> {
+  const qualified = `${quoteIdentifier(schema)}.${quoteIdentifier(table.name)}`;
+  const [indexRows, constraintRows, createRows] = await Promise.all([
+    driver.execute(query(`SHOW INDEXES FROM ${qualified}`)),
+    driver.execute(query(`SHOW CONSTRAINTS FROM ${qualified}`)),
+    driver.execute(query(`SHOW CREATE TABLE ${qualified}`)),
+  ]);
+  const primaryIndexes = new Set(
+    constraintRows
+      .filter((row, index) => textField(row, 'constraint_type', 'cockroach SHOW CONSTRAINTS', index) === 'PRIMARY KEY')
+      .map((row, index) => textField(row, 'constraint_name', 'cockroach SHOW CONSTRAINTS', index)),
+  );
+  const createRow = createRows[0];
+  if (createRow === undefined) throw new Error(`cockroach SHOW CREATE returned no row for ${qualified}`);
+  const predicates = indexPredicates(textField(createRow, 'create_statement', 'cockroach SHOW CREATE TABLE', 0));
+  const columns = table.columns.map(cockroachColumn);
+  return {
+    ...table,
+    columns,
+    indexes: indexes(
+      indexRows.map(parseIndexRow),
+      primaryIndexes,
+      new Set(columns.map(column => column.name)),
+      predicates,
+    ),
+  };
+}
+
+async function cockroachSnapshot(
+  parent: (driver: IntrospectionDriver, options?: IntrospectOptions) => Promise<CatalogSchemaSnapshot>,
+  driver: IntrospectionDriver,
+  options?: IntrospectOptions,
+): Promise<CatalogSchemaSnapshot> {
+  const base = await parent(driver, options);
+  const schemaByTable = await schemasByTable(driver, options);
+  const tables: CatalogTableSnapshot[] = [];
+  for (const table of base.tables) {
+    const schema = schemaByTable.get(table.name);
+    if (schema === undefined)
+      throw new Error(`cockroach introspection could not resolve the schema for "${table.name}"`);
+    tables.push(await cockroachTable(driver, schema, table));
+  }
+  return { ...base, tables };
+}
+
+export const cockroachIntrospector = postgresFamilyIntrospector('cockroach', {
+  snapshot: cockroachSnapshot,
+});

package/src/migrations.ts

@@ -0,0 +1,145 @@
+import type { ChangeOp, ColumnSnapshot, ExtensionType, SchemaSnapshot } from '@zmdb/migrations';
+import { postgresFamilyMigrations } from '@zmdb/postgres';
+import {
+  UnsupportedFeatureError,
+  type MigrationConnection,
+  type MigrationDialect,
+  type MigrationDriver,
+  type MigrationPlan,
+  type MigrationTableOptions,
+  type SchemaObjectOperation,
+} from '@zmdb/sql';
+import { type IndexColumn } from '@zmdb/sql/schema-objects';
+
+export const COCKROACH_TYPE_OVERRIDES = Object.freeze({
+  serial: 'INT8 DEFAULT unique_rowid()',
+  integer: 'INT4',
+});
+
+const parent = postgresFamilyMigrations('cockroach', {
+  types: COCKROACH_TYPE_OVERRIDES,
+});
+
+function extensionName(type: ExtensionType): string {
+  return `${type.extension}.${type.name}`;
+}
+
+function refuseExtension(feature: string): never {
+  throw new UnsupportedFeatureError(
+    feature,
+    'cockroach',
+    `cockroach does not expose PostgreSQL extension installation or extension-backed column types (${feature})`,
+  );
+}
+
+function validateColumn(column: ColumnSnapshot): void {
+  if (typeof column.type !== 'string') refuseExtension(`extension type ${extensionName(column.type)}`);
+}
+
+function validateSnapshot(snapshot: SchemaSnapshot): void {
+  const extension = snapshot.extensions[0];
+  if (extension !== undefined) refuseExtension(`extension "${extension.name}"`);
+  for (const table of snapshot.tables) {
+    for (const column of table.columns) validateColumn(column);
+  }
+}
+
+function validateOperation(operation: ChangeOp): void {
+  switch (operation.kind) {
+    case 'create_extension':
+      refuseExtension(`extension "${operation.name}"`);
+    case 'create_table':
+      for (const column of operation.columns) validateColumn(column);
+      return;
+    case 'add_column':
+      validateColumn(operation.column);
+      return;
+    case 'alter_column_type':
+      if (typeof operation.from !== 'string') refuseExtension(`extension type ${extensionName(operation.from)}`);
+      if (typeof operation.to !== 'string') refuseExtension(`extension type ${extensionName(operation.to)}`);
+      return;
+    default:
+      return;
+  }
+}
+
+function opclass(column: IndexColumn): string | undefined {
+  return typeof column === 'string' ? undefined : column.opclass;
+}
+
+function validateSchemaObject(operation: SchemaObjectOperation): void {
+  switch (operation.kind) {
+    case 'enable_rls':
+    case 'create_policy':
+      throw new UnsupportedFeatureError(
+        'row-level security',
+        'cockroach',
+        'cockroach row-level-security support varies by server release; @zmdb/cockroach refuses PostgreSQL policy DDL',
+      );
+    case 'create_extension':
+      refuseExtension(`extension "${operation.definition.name}"`);
+    case 'create_index': {
+      const method = operation.definition.method;
+      if (method !== undefined) {
+        throw new UnsupportedFeatureError(
+          `index method ${method}`,
+          'cockroach',
+          `cockroach indexes are emitted without a PostgreSQL USING method ("${operation.definition.name}")`,
+        );
+      }
+      const column = operation.definition.columns.find(value => opclass(value) !== undefined);
+      const selectedOpclass = column === undefined ? undefined : opclass(column);
+      if (selectedOpclass !== undefined) {
+        throw new UnsupportedFeatureError(
+          `index operator class ${selectedOpclass}`,
+          'cockroach',
+          `cockroach indexes do not accept PostgreSQL operator classes ("${operation.definition.name}")`,
+        );
+      }
+      return;
+    }
+    default:
+      return;
+  }
+}
+
+export const cockroachMigrations: MigrationDialect<'cockroach'> = Object.freeze({
+  name: 'cockroach',
+  foreignKeyMode: 'deferred',
+  embedded: false,
+  validateSnapshot(snapshot: SchemaSnapshot): void {
+    parent.validateSnapshot(snapshot);
+    validateSnapshot(snapshot);
+  },
+  validatePlan(plan: MigrationPlan): void {
+    parent.validatePlan(plan);
+    validateSnapshot(plan.before);
+    validateSnapshot(plan.after);
+    for (const operation of plan.operations) validateOperation(operation);
+  },
+  ddlType(column: ColumnSnapshot): string {
+    validateColumn(column);
+    return parent.ddlType(column);
+  },
+  emitUp(operation: ChangeOp): string {
+    validateOperation(operation);
+    return parent.emitUp(operation);
+  },
+  emitDown(operation: ChangeOp): string {
+    validateOperation(operation);
+    return parent.emitDown(operation);
+  },
+  emitSchemaObject(operation: SchemaObjectOperation): readonly string[] {
+    validateSchemaObject(operation);
+    return parent.emitSchemaObject(operation);
+  },
+  connection(driver: MigrationDriver<'cockroach'>, options?: MigrationTableOptions): MigrationConnection<'cockroach'> {
+    const connection = parent.connection(driver, options);
+    const { transaction: ignoredTransaction, ...nonTransactional } = connection;
+    void ignoredTransaction;
+    return Object.freeze({
+      ...nonTransactional,
+      transactionalDdl: false,
+    });
+  },
+});