import { getLanguageByResourceType, ColumnIdentity, type ColumnDef, type TableMetadata } from './apps/components/display/dbtable/utils' import { runScriptAndPollResult } from './jobs/utils' import { writingJobOptions } from './jobs/writingJob' import type { DBSchema, SQLSchema } from '$lib/stores' import { stringifySchema } from './copilot/lib' import { datatableReference, type DbInput, type DbType } from './dbTypes' import { withMigrationRole } from './datatableMigrationRole' import { assert } from '$lib/utils' import { WorkspaceService } from '$lib/gen' import { pendingMigrations } from './workspaceSettings/datatableMigrationUtils' import { buildTableEditorValues, type TableEditorForeignKey, type TableEditorValues } from './apps/components/display/dbtable/tableEditor' import { type AlterTableValues } from './apps/components/display/dbtable/queries/alterTable' import { transformForeignKeys, transformSnowflakeForeignKeys, type RawForeignKey } from './apps/components/display/dbtable/queries/relationalKeys' import { groupForeignKeyRows, type DbRelation, type RawAllForeignKeyRow } from './dbRelations' import { joinedColumnDef, joinsPayload, type DbTableJoin } from './dbTableJoins' export type IDbTableOps = { dbType: DbType tableKey: string colDefs: ColumnDef[] getRows: (params: { offset: number limit: number quicksearch: string order_by: string is_desc: boolean /** Raw SQL predicate AND-ed into this read (already escaped). */ whereClause?: string /** The same filters as `whereClause`, column → value, for a source that filters rows * itself. */ columnFilters?: Record /** Whether `order_by` was picked by the user rather than defaulted to the first column. */ explicitSort?: boolean /** Columns of referenced tables to read beside the table's own. */ joins?: DbTableJoin[] }) => Promise getCount: (params: { quicksearch: string whereClause?: string columnFilters?: Record joins?: DbTableJoin[] }) => Promise onUpdate?: ( row: { values: object }, colDef: { field: string; datatype: string }, newValue: string ) => Promise onDelete?: (row: { values: object }) => Promise onInsert?: (row: { values: object }) => Promise } export function dbTableOpsWithPreviewScripts({ input, tableKey, colDefs, workspace, whereClause, version, tag }: { input: DbInput tableKey: string colDefs: ColumnDef[] workspace: string // Optional raw SQL predicate AND-ed into the read queries (count + rows). // Caller-trusted — build it with escaped values. whereClause?: string // DuckLake time-travel: when set, reads are pinned to this catalog snapshot // via `AT (VERSION => n)` (DuckDB/ducklake only). Read-only by nature. version?: number // Overrides the language's native worker tag. tag?: string }): IDbTableOps { const dbType = getDbType(input) const language = getLanguageByResourceType(dbType) // Built per call: an invalid role throws there, as that operation's error, rather than while // the manager renders. const dbArg = () => getDatabaseArg(input) const ducklake = input.type === 'ducklake' ? input.ducklake : undefined function makeMarker(op: string, payload: Record): string { if (ducklake) payload.ducklake = ducklake return `-- WM_INTERNAL_DB_${op} ${JSON.stringify(payload)}` } // A joined column is read as one more column of the table, named by its alias. function readColumns(joins: DbTableJoin[] | undefined) { if (!joins?.length) return { columnDefs: colDefs } return { columnDefs: [...colDefs, ...joins.map(joinedColumnDef)], joins: joinsPayload(joins) } } function combinedWhere(extra: string | undefined): string | undefined { if (!whereClause || !extra) return whereClause || extra || undefined return `(${whereClause}) AND (${extra})` } return { dbType, tableKey, colDefs, getCount: async ({ quicksearch, whereClause: extraWhere, joins }) => { const where = combinedWhere(extraWhere) const content = makeMarker('COUNT', { table: tableKey, ...readColumns(joins), ...(where ? { whereClause: where } : {}), ...(version != undefined ? { version } : {}) }) const result = await runScriptAndPollResult({ workspace, requestBody: { args: { ...dbArg(), quicksearch }, language, content, tag } }) const count = result?.[0].count as number return count }, getRows: async ({ whereClause: extraWhere, columnFilters: _, explicitSort: __, joins, ...params }) => { const where = combinedWhere(extraWhere) const content = makeMarker('SELECT', { table: tableKey, ...readColumns(joins), fixPgIntTypes: true, ...(where ? { whereClause: where } : {}), ...(version != undefined ? { version } : {}) }) let items = (await runScriptAndPollResult({ workspace, requestBody: { args: { ...dbArg(), ...params }, language, content, tag } })) as unknown[] if (!items || !Array.isArray(items)) { throw 'items is not an array' } return items }, onUpdate: async ({ values }, colDef, newValue) => { const content = makeMarker('UPDATE', { table: tableKey, column: colDef, columns: colDefs }) await runScriptAndPollResult( { workspace, requestBody: { args: { ...dbArg(), value_to_update: newValue, ...values }, language, content, tag } }, writingJobOptions ) }, onDelete: async ({ values }) => { const content = makeMarker('DELETE', { table: tableKey, columns: colDefs }) await runScriptAndPollResult( { workspace, requestBody: { args: { ...dbArg(), ...values }, language, content, tag } }, writingJobOptions ) }, onInsert: async ({ values }) => { const content = makeMarker('INSERT', { table: tableKey, columns: colDefs }) await runScriptAndPollResult( { workspace, requestBody: { args: { ...dbArg(), ...values }, language, content, tag } }, writingJobOptions ) } } } export type DucklakeSnapshot = { snapshot_id: number // DuckLake returns this as microseconds-since-epoch serialized as a string // (TIMESTAMP); callers must convert before formatting. snapshot_time: string | number } /** * Column metadata of a ducklake table *at a specific snapshot*. The catalog's * `information_schema` only reflects the current schema, so a time-travel read * pinned to an older version must enumerate the columns that existed *then* — * otherwise a column added in a later snapshot would break the `SELECT … AT * (VERSION => n)`. `DESCRIBE SELECT * FROM … AT (VERSION => n)` gives exactly * that. Returns minimal `ColumnDef`s (field + datatype) — enough for the * read-only preview's SELECT/COUNT and grid headers. */ export async function fetchDucklakeColumnsAtVersion({ workspace, ducklake, tableKey, version }: { workspace: string ducklake: string tableKey: string version: number }): Promise { // Quote each identifier part (schema.table) so a dotted/odd table name can't // break the statement, and double single-quotes in the catalog name so it // can't break out of the ATTACH string literal (mirrors the backend's // `escape_sql_literal`). `version` is a number — injection-safe. const quoted = tableKey .split('.') .map((p) => `"${p.replace(/"/g, '""')}"`) .join('.') const ducklakeLit = ducklake.replace(/'/g, "''") const content = `ATTACH 'ducklake://${ducklakeLit}' AS __dlv__; USE __dlv__; ` + `DESCRIBE SELECT * FROM ${quoted} AT (VERSION => ${version});` const rows = (await runScriptAndPollResult({ workspace, requestBody: { args: {}, language: 'duckdb', content } })) as { column_name: string; column_type: string }[] if (!Array.isArray(rows)) return [] return rows.map((r) => ({ field: r.column_name, datatype: r.column_type, defaultvalue: '', isprimarykey: false, isidentity: ColumnIdentity.No, isnullable: 'YES' as const, isenum: false })) } /** * List a ducklake table's time-travel history, newest first. DuckLake snapshots * are catalog-wide commits; passing `table` (schema-qualified, e.g. * `main.events_daily`) scopes the list to snapshots where the table exists — * otherwise an `AT (VERSION => n)` read could target a version predating the * table's creation and error. Runs the `DUCKLAKE_SNAPSHOTS` marker as a duckdb * preview job (server-side SQL build + ATTACH), so no raw SQL is constructed in * the client. */ export async function fetchDucklakeSnapshots({ workspace, ducklake, table }: { workspace: string ducklake: string table?: string }): Promise { const content = `-- WM_INTERNAL_DB_DUCKLAKE_SNAPSHOTS ${JSON.stringify({ ducklake, ...(table ? { table } : {}) })}` const rows = await runScriptAndPollResult({ workspace, requestBody: { args: {}, language: 'duckdb', content } }) return Array.isArray(rows) ? (rows as DucklakeSnapshot[]) : [] } export type IDbSchemaOps = { onDelete: (params: { tableKey: string; schema?: string }) => Promise onCreate: (params: { values: TableEditorValues; schema?: string }) => Promise previewCreateSql: (params: { values: TableEditorValues; schema?: string }) => Promise onAlter: (params: { values: AlterTableValues /** Reverse diff (new → old), used to generate the down migration. */ reverse?: AlterTableValues schema?: string }) => Promise previewAlterSql: (params: { values: AlterTableValues; schema?: string }) => Promise onCreateSchema: (params: { schema: string }) => Promise onDeleteSchema: (params: { schema: string }) => Promise onRenameSchema: (params: { schema: string; newSchema: string }) => Promise onFetchTableEditorDefinition: (params: { table: string schema?: string colDefs: TableMetadata }) => Promise onFetchForeignKeys: (params: { table: string schema?: string }) => Promise /** Every foreign key of the database in one query, for the schema diagram. * PostgreSQL only — the other databases either don't expose foreign keys or * can only be asked one table at a time. */ onFetchAllForeignKeys: () => Promise } /** Thrown by a schema op when the user declines the out-of-order run warning. * Callers should treat it as a silent cancel (no error toast). */ export class MigrationRunCancelled extends Error { constructor() { super('Migration run cancelled') this.name = 'MigrationRunCancelled' } } export function dbSchemaOpsWithPreviewScripts({ workspace, input, confirmRunOutOfOrder, tag }: { workspace: string input: DbInput /** Asked before running a just-created migration ahead of `pendingCount` * still-pending earlier ones. Return false to abort (throws MigrationRunCancelled). */ confirmRunOutOfOrder?: (pendingCount: number) => Promise // Overrides the language's native worker tag, for a database only reachable from // a specific worker group. Data table migrations are run by the server, which // picks their tag itself, so they are unaffected. tag?: string }): IDbSchemaOps { const dbType = getDbType(input) // Built per call, for the same reason as in the table ops above. const dbArg = () => getDatabaseArg(input) const language = getLanguageByResourceType(dbType) const ducklake = input.type === 'ducklake' ? input.ducklake : undefined // When managing a data table, schema changes are recorded as migrations // instead of being run ad-hoc, so the manager stays the source of truth. const datatableName = input.type === 'database' && input.resourcePath.startsWith('datatable://') ? input.resourcePath.slice('datatable://'.length) : undefined // A migration declaring no role runs as admin, whatever role the manager connects as. const migrationRole = input.type === 'database' ? (input.role ?? input.migrationRole) : undefined function makeMarker(op: string, payload: Record): string { if (ducklake) payload.ducklake = ducklake return `-- WM_INTERNAL_DB_${op} ${JSON.stringify(payload)}` } // Auto-generated migration name, e.g. `create_customers`. The server allocates // a unique timestamp (bumping on collision), so the name itself need not be unique. function migrationName(op: string, target: string): string { const safe = target.replace(/[^a-zA-Z0-9_-]+/g, '_').replace(/^_+|_+$/g, '') return safe ? `${op}_${safe}` : op } // A dropped SERIAL column reports its default as `nextval()` of an owned // sequence that is dropped along with the column. When the down re-adds such a // column, recreate it as its SERIAL type so a fresh sequence is created // instead of referencing the gone one. const SERIAL_FOR: Record = { BIGINT: 'BIGSERIAL', INT8: 'BIGSERIAL', INTEGER: 'SERIAL', INT: 'SERIAL', INT4: 'SERIAL', SMALLINT: 'SMALLSERIAL', INT2: 'SMALLSERIAL' } function reverseSerialFix(reverse: AlterTableValues): AlterTableValues { return { ...reverse, operations: reverse.operations.map((op) => { if (op.kind !== 'addColumn' || !/nextval\s*\(/i.test(op.column.defaultValue ?? '')) { return op } const serial = SERIAL_FOR[(op.column.datatype ?? '').toUpperCase()] return serial ? { ...op, column: { ...op.column, datatype: serial, defaultValue: undefined } } : op }) } } // Frame a single (or multi-) statement body in an explicit, `;`-terminated // transaction, matching the data table migration convention. Some expanded // markers (e.g. ALTER TABLE) already come wrapped in their own transaction, so // avoid nesting BEGIN/COMMIT in that case. function wrapMigration(sql: string): string { const t = sql.trim() if (/^BEGIN\b/i.test(t)) return t return `BEGIN;\n\n${t.endsWith(';') ? t : `${t};`}\n\nEND;` } // Apply a DDL marker. For a data table that has migrations enabled this // creates a migration and runs it (rolling the record back if the run fails); // otherwise it runs ad-hoc via the internal-db job as before. `downContent`, // when provided, is expanded into the migration's down SQL (Postgres only). async function applyDdl(migName: string, content: string, downContent?: string): Promise { // A DDL edit on a migrations-enabled data table must be captured as a // migration. Don't swallow a status-check failure by defaulting to ad-hoc: // that would run the change untracked (schema drift) — exactly what this // feature prevents. Let the error propagate (fail closed); only fall back to // ad-hoc when there's no data table, or `enabled === false` is returned. const status = datatableName ? await WorkspaceService.getDatatableMigrationsStatus({ workspace, datatableName }) : undefined if (!datatableName || !status?.enabled) { await runScriptAndPollResult( { workspace, requestBody: { args: dbArg(), content, language, tag } }, writingJobOptions ) return } // The new migration gets the highest timestamp, so any still-pending // migration is earlier: running only this one applies it out of order. // Warn like the row-level Run action does (skipped if no confirm hook). if (confirmRunOutOfOrder) { const pending = pendingMigrations(status.migrations).length if (pending > 0 && !(await confirmRunOutOfOrder(pending))) { throw new MigrationRunCancelled() } } // Wrapped before annotating: the annotation must lead, above `BEGIN;`. const codeUp = withMigrationRole( wrapMigration(await expandMarker(workspace, language, content)), migrationRole ) // Down migrations are only generated for Postgres for now. let codeDown: string | undefined if (downContent && dbType === 'postgresql') { const downSql = (await expandMarker(workspace, language, downContent)).trim() if (downSql) codeDown = withMigrationRole(wrapMigration(downSql), migrationRole) } const created = await WorkspaceService.createDatatableMigration({ workspace, datatableName, requestBody: { name: migName, code_up: codeUp, ...(codeDown ? { code_down: codeDown } : {}) } }) try { await WorkspaceService.runDatatableMigrations({ workspace, datatableName, only: created.timestamp }) } catch (e) { await WorkspaceService.deleteDatatableMigration({ workspace, datatableName, timestamp: created.timestamp }).catch(() => {}) throw e } } /** Resolves to [] when the database has no foreign key introspection * (BigQuery) or the query fails: callers treat foreign keys as optional. */ async function fetchForeignKeys({ table, schema }: { table: string schema?: string }): Promise { if (dbType === 'bigquery') return [] try { const fkContent = makeMarker('FOREIGN_KEYS', { table, schema }) const fkResult = await runScriptAndPollResult({ workspace, requestBody: { args: dbArg(), content: fkContent, language, tag } }) let rawForeignKeys: RawForeignKey[] if (dbType === 'snowflake') { rawForeignKeys = transformSnowflakeForeignKeys(fkResult as any[]) } else { rawForeignKeys = fkResult as RawForeignKey[] if (rawForeignKeys && Array.isArray(rawForeignKeys)) { rawForeignKeys = rawForeignKeys.map((fk) => { const lowerFk: any = {} Object.keys(fk).forEach((key) => { lowerFk[key.toLowerCase()] = fk[key] }) return lowerFk }) } } if (rawForeignKeys && Array.isArray(rawForeignKeys)) { return transformForeignKeys(rawForeignKeys) } } catch (e) { console.warn('Failed to fetch foreign keys:', e) } return [] } async function fetchAllForeignKeys(): Promise { if (dbType !== 'postgresql') return [] const content = makeMarker('ALL_FOREIGN_KEYS', {}) const rows = (await runScriptAndPollResult({ workspace, requestBody: { args: dbArg(), content, language, tag } })) as RawAllForeignKeyRow[] return Array.isArray(rows) ? groupForeignKeyRows(rows) : [] } return { onDelete: async ({ tableKey, schema }) => { const content = makeMarker('DROP_TABLE', { table: tableKey, schema }) await applyDdl(migrationName('drop', tableKey), content) }, onCreate: async ({ values, schema }) => { const content = makeMarker('CREATE_TABLE', { name: values.name, columns: values.columns, foreignKeys: values.foreignKeys, schema }) const downContent = makeMarker('DROP_TABLE', { table: values.name, schema }) await applyDdl(migrationName('create', values.name), content, downContent) }, previewCreateSql: async ({ values, schema }) => { const content = makeMarker('CREATE_TABLE', { name: values.name, columns: values.columns, foreignKeys: values.foreignKeys, schema }) return expandMarker(workspace, language, content) }, onAlter: async ({ values, reverse, schema }) => { const content = makeMarker('ALTER_TABLE', { name: values.name, operations: values.operations, schema }) // The down is the same alter run in the opposite direction. const downContent = reverse ? makeMarker('ALTER_TABLE', { name: reverse.name, operations: reverseSerialFix(reverse).operations, schema }) : undefined await applyDdl(migrationName('alter', values.name), content, downContent) }, previewAlterSql: async ({ values, schema }) => { const content = makeMarker('ALTER_TABLE', { name: values.name, operations: values.operations, schema }) return expandMarker(workspace, language, content) }, onCreateSchema: async ({ schema }) => { const content = makeMarker('CREATE_SCHEMA', { schema }) const downContent = makeMarker('DROP_SCHEMA', { schema }) await applyDdl(migrationName('create_schema', schema), content, downContent) }, onDeleteSchema: async ({ schema }) => { const content = makeMarker('DROP_SCHEMA', { schema }) const downContent = makeMarker('CREATE_SCHEMA', { schema }) await applyDdl(migrationName('drop_schema', schema), content, downContent) }, onRenameSchema: async ({ schema, newSchema }) => { const content = makeMarker('RENAME_SCHEMA', { schema, new_schema: newSchema }) const downContent = makeMarker('RENAME_SCHEMA', { schema: newSchema, new_schema: schema }) await applyDdl(migrationName('rename_schema', schema), content, downContent) }, onFetchForeignKeys: fetchForeignKeys, onFetchAllForeignKeys: fetchAllForeignKeys, onFetchTableEditorDefinition: async ({ table, schema, colDefs }) => { const foreignKeys = await fetchForeignKeys({ table, schema }) let pk_constraint_name: string | undefined // Fetch primary key constraint name (not supported for BigQuery/MySQL) if (dbType !== 'bigquery' && dbType !== 'mysql') { try { const pkContent = makeMarker('PRIMARY_KEY_CONSTRAINT', { table, schema }) const pkResult = (await runScriptAndPollResult({ workspace, requestBody: { args: dbArg(), content: pkContent, language, tag } })) as { constraint_name?: string; CONSTRAINT_NAME?: string }[] if (pkResult && Array.isArray(pkResult) && pkResult.length > 0) { const pkRecord: any = pkResult[0] pk_constraint_name = pkRecord?.constraint_name || pkRecord?.CONSTRAINT_NAME || '' } } catch (e) { console.warn('Failed to fetch primary key constraint:', e) } } return buildTableEditorValues({ tableName: table, metadata: colDefs, foreignKeys, pk_constraint_name }) } } } export async function getDucklakeSchema({ workspace, ducklake, tag }: { workspace: string ducklake: string tag?: string }): Promise { let result = await runScriptAndPollResult({ workspace, requestBody: { language: 'duckdb', content: `ATTACH 'ducklake://${ducklake}' AS __ducklake__; ${DUCKLAKE_GET_SCHEMA_QUERY}`, args: {}, tag } }) let schemas = Array.isArray(result) && result.length && (result?.[0]?.['result'] ?? {}) // Safety for agent workers (duckdb ffi lib used to return JSON as stringified json) if (typeof schemas === 'string') schemas = JSON.parse(schemas) if (!schemas) throw new Error('Failed to get Ducklake schema: ' + JSON.stringify(result)) assert('schemas is an object', typeof schemas === 'object') let schema: Omit = { schema: schemas, publicOnly: false, lang: 'ducklake' } return { ...schema, stringified: stringifySchema(schema) } } // Returns every schema in the ducklake (including empty ones, e.g. freshly created) // as a nested map { schema: { table: { column: {...} } } }. const DUCKLAKE_GET_SCHEMA_QUERY = ` SELECT json_group_object(schema_name, COALESCE(schema_data, json_object())) AS result FROM ( SELECT s.schema_name, ( SELECT json_group_object(table_name, table_data) FROM ( SELECT c.table_name, json_group_object( c.column_name, json_object( 'type', c.data_type, 'default', c.column_default, 'required', c.is_nullable == 'NO' ) ) AS table_data FROM information_schema.columns c WHERE c.table_catalog = '__ducklake__' AND c.table_schema = s.schema_name GROUP BY c.table_name ) ) AS schema_data FROM information_schema.schemata s WHERE s.catalog_name = '__ducklake__' )` export function getDbType(input: DbInput): DbType { switch (input.type) { case 'database': return input.resourceType case 'ducklake': return 'duckdb' } } /** Tag the database's jobs carry with no override: for every DB the manager * supports, the language of its queries is also the name of its native tag. */ export function getDefaultDbTag(input: DbInput): string { return getLanguageByResourceType(getDbType(input)) } export function getDatabaseArg(input: DbInput | undefined) { if (input?.type === 'database') { if (input.resourcePath.startsWith('datatable://')) { return { database: datatableReference(input.resourcePath.slice('datatable://'.length), input.role) } } else { return { database: '$res:' + input.resourcePath } } } return {} } async function expandMarker(workspace: string, language: string, content: string): Promise { const response = await fetch(`/api/w/${workspace}/internal_db/expand_marker`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ language, content }) }) if (!response.ok) { throw new Error(await response.text()) } const result = (await response.json()) as { code: string } return result.code }