Files
ThothII/backend/src/catalog/schema-introspector.ts
T

499 lines
21 KiB
TypeScript

import { readFile } from "node:fs/promises";
import { z } from "zod";
import type { WorkspaceSecretStore } from "../workspaces/secret-store.js";
import type { CatalogPostgresAccess } from "./postgres-access.js";
import { CATALOG_SECRET_IDS } from "./secrets.js";
import {
CatalogConnectorError,
CatalogSchemaCapabilityUnavailableError,
type CatalogSyncPhase,
type ObservedCatalogColumn,
type ObservedCatalogRelationship,
type ObservedCatalogTable,
type ObservedSchemaSnapshot,
type WorkspaceDatabase,
} from "./types.js";
export type CatalogSchemaScanProgress = (
phase: Extract<CatalogSyncPhase, "connecting" | "scanning_tables" | "scanning_columns" | "scanning_relationships">,
counts?: { tables?: number; columns?: number; relationships?: number },
) => Promise<void> | void;
export interface CatalogSchemaIntrospector {
scan(
database: WorkspaceDatabase,
signal: AbortSignal,
progress?: CatalogSchemaScanProgress,
): Promise<ObservedSchemaSnapshot>;
}
const identifier = z.string().min(1).max(128);
const nullableText = z.string().nullable();
const capability = z.enum(["available", "unavailable"]);
const restSnapshotSchema = z.object({
schemaVersion: z.literal(1),
capabilities: z.object({
tables: capability,
columns: capability,
relationships: capability,
}).strict(),
tables: z.array(z.object({
name: identifier,
sourceComment: nullableText,
}).strict()),
columns: z.array(z.object({
tableName: identifier,
name: identifier,
ordinalPosition: z.number().int().positive(),
dataType: z.string().min(1).max(2_000),
isNullable: z.boolean(),
defaultExpression: nullableText,
primaryKeyPosition: z.number().int().positive().nullable(),
sourceComment: nullableText,
}).strict()),
relationships: z.array(z.object({
constraintName: identifier,
sourceTableName: identifier,
targetTableName: identifier,
updateRule: z.string().min(1).max(64),
deleteRule: z.string().min(1).max(64),
deferrable: z.boolean(),
initiallyDeferred: z.boolean(),
columns: z.array(z.object({
position: z.number().int().positive(),
sourceColumnName: identifier,
targetColumnName: identifier,
}).strict()).min(1),
}).strict()),
}).strict();
function sqlString(value: string): string {
return `'${value.replaceAll("'", "''")}'`;
}
function restSnapshotQuery(schemaName: string): string {
const schema = sqlString(schemaName);
return `WITH target_schema AS (
SELECT oid
FROM pg_catalog.pg_namespace
WHERE nspname = ${schema}
),
observed_tables AS (
SELECT c.oid,
c.relname AS name,
d.description AS source_comment
FROM pg_catalog.pg_class c
JOIN target_schema n ON n.oid = c.relnamespace
LEFT JOIN pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = 0
WHERE c.relkind IN ('r', 'p')
),
primary_key_columns AS (
SELECT i.indrelid AS table_oid,
key.attnum,
key.ordinality::integer AS position
FROM pg_catalog.pg_index i
CROSS JOIN LATERAL unnest(i.indkey) WITH ORDINALITY AS key(attnum, ordinality)
WHERE i.indisprimary
),
observed_columns AS (
SELECT table_info.name AS table_name,
a.attname AS name,
a.attnum::integer AS ordinal_position,
pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
NOT a.attnotnull AS is_nullable,
pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) AS default_expression,
pk.position AS primary_key_position,
d.description AS source_comment
FROM observed_tables table_info
JOIN pg_catalog.pg_attribute a ON a.attrelid = table_info.oid
LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = table_info.oid AND ad.adnum = a.attnum
LEFT JOIN pg_catalog.pg_description d ON d.objoid = table_info.oid AND d.objsubid = a.attnum
LEFT JOIN primary_key_columns pk ON pk.table_oid = table_info.oid AND pk.attnum = a.attnum
WHERE a.attnum > 0
AND NOT a.attisdropped
),
relationship_pairs AS (
SELECT con.oid AS constraint_oid,
con.conname AS constraint_name,
source_table.relname AS source_table_name,
target_table.relname AS target_table_name,
CASE con.confupdtype
WHEN 'a' THEN 'NO ACTION'
WHEN 'r' THEN 'RESTRICT'
WHEN 'c' THEN 'CASCADE'
WHEN 'n' THEN 'SET NULL'
WHEN 'd' THEN 'SET DEFAULT'
END AS update_rule,
CASE con.confdeltype
WHEN 'a' THEN 'NO ACTION'
WHEN 'r' THEN 'RESTRICT'
WHEN 'c' THEN 'CASCADE'
WHEN 'n' THEN 'SET NULL'
WHEN 'd' THEN 'SET DEFAULT'
END AS delete_rule,
con.condeferrable AS is_deferrable,
con.condeferred AS initially_deferred,
source_key.ordinality::integer AS position,
source_column.attname AS source_column_name,
target_column.attname AS target_column_name
FROM pg_catalog.pg_constraint con
JOIN pg_catalog.pg_class source_table ON source_table.oid = con.conrelid
JOIN target_schema source_namespace ON source_namespace.oid = source_table.relnamespace
JOIN pg_catalog.pg_class target_table ON target_table.oid = con.confrelid
JOIN target_schema target_namespace ON target_namespace.oid = target_table.relnamespace
JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS source_key(attnum, ordinality) ON true
JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS target_key(attnum, ordinality)
ON target_key.ordinality = source_key.ordinality
JOIN pg_catalog.pg_attribute source_column
ON source_column.attrelid = source_table.oid AND source_column.attnum = source_key.attnum
JOIN pg_catalog.pg_attribute target_column
ON target_column.attrelid = target_table.oid AND target_column.attnum = target_key.attnum
WHERE con.contype = 'f'
),
observed_relationships AS (
SELECT constraint_oid,
constraint_name,
source_table_name,
target_table_name,
update_rule,
delete_rule,
is_deferrable,
initially_deferred,
pg_catalog.jsonb_agg(
pg_catalog.jsonb_build_object(
'position', position,
'sourceColumnName', source_column_name,
'targetColumnName', target_column_name
) ORDER BY position
) AS columns
FROM relationship_pairs
GROUP BY constraint_oid, constraint_name, source_table_name, target_table_name,
update_rule, delete_rule, is_deferrable, initially_deferred
)
SELECT 1 AS "schemaVersion",
pg_catalog.jsonb_build_object(
'tables', 'available',
'columns', 'available',
'relationships', 'available'
) AS capabilities,
COALESCE((
SELECT pg_catalog.jsonb_agg(
pg_catalog.jsonb_build_object('name', name, 'sourceComment', source_comment)
ORDER BY name
)
FROM observed_tables
), '[]'::jsonb) AS tables,
COALESCE((
SELECT pg_catalog.jsonb_agg(
pg_catalog.jsonb_build_object(
'tableName', table_name,
'name', name,
'ordinalPosition', ordinal_position,
'dataType', data_type,
'isNullable', is_nullable,
'defaultExpression', default_expression,
'primaryKeyPosition', primary_key_position,
'sourceComment', source_comment
) ORDER BY table_name, ordinal_position
)
FROM observed_columns
), '[]'::jsonb) AS columns,
COALESCE((
SELECT pg_catalog.jsonb_agg(
pg_catalog.jsonb_build_object(
'constraintName', constraint_name,
'sourceTableName', source_table_name,
'targetTableName', target_table_name,
'updateRule', update_rule,
'deleteRule', delete_rule,
'deferrable', is_deferrable,
'initiallyDeferred', initially_deferred,
'columns', columns
) ORDER BY source_table_name, constraint_name
)
FROM observed_relationships
), '[]'::jsonb) AS relationships
FROM target_schema`;
}
function required(value: string | undefined): string {
if (!value) throw new CatalogConnectorError("Database binding is incomplete");
return value;
}
function textOrNull(value: unknown): string | null {
return typeof value === "string" && value.length > 0 ? value : null;
}
function actionRule(value: unknown): string {
const rules: Record<string, string> = {
a: "NO ACTION",
r: "RESTRICT",
c: "CASCADE",
n: "SET NULL",
d: "SET DEFAULT",
};
const rule = rules[String(value)];
if (!rule) throw new CatalogConnectorError("Schema introspection returned an unknown relationship action");
return rule;
}
function normalized(snapshot: ObservedSchemaSnapshot): ObservedSchemaSnapshot {
const tables = new Map<string, ObservedCatalogTable>();
for (const table of snapshot.tables) {
if (tables.has(table.name)) throw new CatalogConnectorError("Schema introspection returned duplicate tables");
tables.set(table.name, table);
}
const columns = new Map<string, ObservedCatalogColumn>();
for (const column of snapshot.columns) {
const key = `${column.tableName}\u0000${column.name}`;
if (columns.has(key)) throw new CatalogConnectorError("Schema introspection returned duplicate columns");
columns.set(key, column);
}
const relationships = new Map<string, ObservedCatalogRelationship>();
for (const relationship of snapshot.relationships) {
const key = `${relationship.sourceTableName}\u0000${relationship.constraintName}`;
if (relationships.has(key)) throw new CatalogConnectorError("Schema introspection returned duplicate relationships");
relationships.set(key, {
...relationship,
columns: [...relationship.columns].sort((a, b) => a.position - b.position),
});
}
return {
schemaVersion: 1,
capabilities: snapshot.capabilities,
tables: [...tables.values()].sort((a, b) => a.name.localeCompare(b.name)),
columns: [...columns.values()].sort((a, b) => (
a.tableName.localeCompare(b.tableName) || a.ordinalPosition - b.ordinalPosition
)),
relationships: [...relationships.values()].sort((a, b) => (
a.sourceTableName.localeCompare(b.sourceTableName) || a.constraintName.localeCompare(b.constraintName)
)),
};
}
export class ConcreteCatalogSchemaIntrospector implements CatalogSchemaIntrospector {
constructor(
private readonly postgres: CatalogPostgresAccess,
private readonly secretStore: WorkspaceSecretStore,
) {}
async scan(
database: WorkspaceDatabase,
signal: AbortSignal,
progress?: CatalogSchemaScanProgress,
): Promise<ObservedSchemaSnapshot> {
return database.binding.transport === "rest_api"
? await this.scanRest(database, signal, progress)
: await this.scanPostgres(database, signal, progress);
}
private async scanPostgres(
database: WorkspaceDatabase,
signal: AbortSignal,
progress?: CatalogSchemaScanProgress,
): Promise<ObservedSchemaSnapshot> {
await progress?.("connecting");
const client = await this.postgres.connect(database, signal);
try {
const schema = await client.query(
"SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_namespace WHERE nspname = $1) AS present",
[database.schema],
);
if (schema.rows[0]?.present !== true) throw new CatalogConnectorError("Database schema is unavailable");
await progress?.("scanning_tables");
const tableResult = await client.query(
`SELECT c.relname AS name, d.description AS source_comment
FROM pg_catalog.pg_class c
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
LEFT JOIN pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = 0
WHERE c.relkind IN ('r', 'p') AND n.nspname = $1
ORDER BY c.relname`,
[database.schema],
);
const tables: ObservedCatalogTable[] = tableResult.rows.map((row) => ({
name: String(row.name),
sourceComment: textOrNull(row.source_comment),
}));
await progress?.("scanning_tables", { tables: tables.length });
await progress?.("scanning_columns", { tables: tables.length });
const columnResult = await client.query(
`SELECT c.relname AS table_name,
a.attname AS name,
a.attnum::integer AS ordinal_position,
pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
NOT a.attnotnull AS is_nullable,
pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) AS default_expression,
pk.position AS primary_key_position,
d.description AS source_comment
FROM pg_catalog.pg_class c
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid
LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = c.oid AND ad.adnum = a.attnum
LEFT JOIN pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
LEFT JOIN LATERAL (
SELECT key.ordinality::integer AS position
FROM pg_catalog.pg_index i
CROSS JOIN LATERAL unnest(i.indkey) WITH ORDINALITY AS key(attnum, ordinality)
WHERE i.indrelid = c.oid AND i.indisprimary AND key.attnum = a.attnum
LIMIT 1
) pk ON true
WHERE c.relkind IN ('r', 'p')
AND n.nspname = $1
AND a.attnum > 0
AND NOT a.attisdropped
ORDER BY c.relname, a.attnum`,
[database.schema],
);
const columns: ObservedCatalogColumn[] = columnResult.rows.map((row) => ({
tableName: String(row.table_name),
name: String(row.name),
ordinalPosition: Number(row.ordinal_position),
dataType: String(row.data_type),
isNullable: row.is_nullable === true,
defaultExpression: textOrNull(row.default_expression),
primaryKeyPosition: row.primary_key_position === null || row.primary_key_position === undefined
? null
: Number(row.primary_key_position),
sourceComment: textOrNull(row.source_comment),
}));
await progress?.("scanning_columns", { tables: tables.length, columns: columns.length });
await progress?.("scanning_relationships", { tables: tables.length, columns: columns.length });
const relationshipResult = await client.query(
`SELECT con.conname AS constraint_name,
source_table.relname AS source_table_name,
target_table.relname AS target_table_name,
con.confupdtype AS update_action,
con.confdeltype AS delete_action,
con.condeferrable AS deferrable,
con.condeferred AS initially_deferred,
source_key.ordinality::integer AS position,
source_column.attname AS source_column_name,
target_column.attname AS target_column_name
FROM pg_catalog.pg_constraint con
JOIN pg_catalog.pg_class source_table ON source_table.oid = con.conrelid
JOIN pg_catalog.pg_namespace source_namespace ON source_namespace.oid = source_table.relnamespace
JOIN pg_catalog.pg_class target_table ON target_table.oid = con.confrelid
JOIN pg_catalog.pg_namespace target_namespace ON target_namespace.oid = target_table.relnamespace
JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS source_key(attnum, ordinality) ON true
JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS target_key(attnum, ordinality)
ON target_key.ordinality = source_key.ordinality
JOIN pg_catalog.pg_attribute source_column
ON source_column.attrelid = source_table.oid AND source_column.attnum = source_key.attnum
JOIN pg_catalog.pg_attribute target_column
ON target_column.attrelid = target_table.oid AND target_column.attnum = target_key.attnum
WHERE con.contype = 'f'
AND source_namespace.nspname = $1
AND target_namespace.nspname = $1
ORDER BY source_table.relname, con.conname, source_key.ordinality`,
[database.schema],
);
const relationshipMap = new Map<string, ObservedCatalogRelationship>();
for (const row of relationshipResult.rows) {
const sourceTableName = String(row.source_table_name);
const constraintName = String(row.constraint_name);
const key = `${sourceTableName}\u0000${constraintName}`;
const current = relationshipMap.get(key) ?? {
constraintName,
sourceTableName,
targetTableName: String(row.target_table_name),
updateRule: actionRule(row.update_action),
deleteRule: actionRule(row.delete_action),
deferrable: row.deferrable === true,
initiallyDeferred: row.initially_deferred === true,
columns: [],
};
current.columns.push({
position: Number(row.position),
sourceColumnName: String(row.source_column_name),
targetColumnName: String(row.target_column_name),
});
relationshipMap.set(key, current);
}
const relationships = [...relationshipMap.values()];
await progress?.("scanning_relationships", {
tables: tables.length,
columns: columns.length,
relationships: relationships.length,
});
return normalized({
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables,
columns,
relationships,
});
} finally {
await client.end();
}
}
private async scanRest(
database: WorkspaceDatabase,
signal: AbortSignal,
progress?: CatalogSchemaScanProgress,
): Promise<ObservedSchemaSnapshot> {
await progress?.("connecting");
const auth = database.binding.restAuth ?? "bearer";
const materialized = this.secretStore.materialize(
database.workspaceId,
auth === "none" ? [] : [CATALOG_SECRET_IDS.apiKey],
);
try {
const headers: Record<string, string> = { "content-type": "application/json" };
if (auth !== "none") {
const credentialFile = materialized.files.get(CATALOG_SECRET_IDS.apiKey);
if (!credentialFile) throw new CatalogConnectorError("REST API key is not configured");
const credential = (await readFile(credentialFile, "utf8")).trim();
if (auth === "bearer") headers.authorization = `Bearer ${credential}`;
else headers["x-api-key"] = credential;
}
const baseUrl = required(database.binding.baseUrl).replace(/\/+$/, "");
const response = await fetch(`${baseUrl}/rpc/schema_snapshot`, {
method: "POST",
headers,
body: JSON.stringify({ schema_name: database.schema }),
signal,
});
let body: unknown;
if (response.ok) {
body = await response.json();
} else if (response.status === 404) {
const fallback = await fetch(`${baseUrl}/rpc/run_query`, {
method: "POST",
headers,
body: JSON.stringify({ query_text: restSnapshotQuery(database.schema) }),
signal,
});
if (!fallback.ok) throw new CatalogSchemaCapabilityUnavailableError("schema_snapshot");
const rows: unknown = await fallback.json();
if (!Array.isArray(rows) || rows.length !== 1) {
throw new CatalogConnectorError("REST schema snapshot fallback is invalid");
}
body = rows[0];
} else {
throw new CatalogSchemaCapabilityUnavailableError("schema_snapshot");
}
const parsed = restSnapshotSchema.safeParse(body);
if (!parsed.success) throw new CatalogConnectorError("REST schema snapshot is invalid");
const snapshot = normalized(parsed.data);
await progress?.("scanning_tables", { tables: snapshot.tables.length });
await progress?.("scanning_columns", { tables: snapshot.tables.length, columns: snapshot.columns.length });
await progress?.("scanning_relationships", {
tables: snapshot.tables.length,
columns: snapshot.columns.length,
relationships: snapshot.relationships.length,
});
return snapshot;
} catch (error) {
if (error instanceof CatalogConnectorError) throw error;
throw new CatalogConnectorError("REST schema introspection failed");
} finally {
materialized.release();
}
}
}