Line data Source code
1 : import * as arrow from "apache-arrow";
2 : import * as dashql from '@ankoh/dashql-core';
3 :
4 : import { QueryExecutor } from '../query_executor.js';
5 : import { QueryExecutionArgs } from '../query_execution_args.js';
6 : import { DynamicConnectionDispatch } from '../connection_registry.js';
7 : import { CATALOG_UPDATE_SCHEMA_SCRIPT, CATALOG_UPDATE_REGISTER_QUERY } from "../connection_state.js";
8 : import { QueryType } from '../../query/query_execution_state.js';
9 : import { generateSchemaSQL, generateCatalogScriptHeader, CatalogSource, quoteIdentifier, type ColumnMetadata } from './catalog_sql_generator.js';
10 : import { generateFunctionScriptHeader } from './catalog_function_sql_generator.js';
11 : import { queryPgProc, generateCatalogSQLFromPgProc } from './catalog_query_pg_proc.js';
12 : import { type LoggerLike } from '../../platform/logger/logger.js';
13 : import { beginCatalogBatch, publishCatalogBatch, type CatalogBatchEntry, type CatalogBatchOptions } from './catalog_batch.js';
14 :
15 1 : const LOG_CTX = "catalog_pg";
16 1 : export const CATALOG_QUERY_READ_TIMEOUT_MS = 60_000;
17 :
18 : export type PgAttributeColumnsTable = arrow.Table<{
19 : table_schema: arrow.Utf8;
20 : table_name: arrow.Utf8;
21 : column_name: arrow.Utf8;
22 : ordinal_position: arrow.Int32;
23 : data_type: arrow.Utf8;
24 : is_nullable: arrow.Utf8;
25 : numeric_precision: arrow.Int32;
26 : numeric_scale: arrow.Int32;
27 : }>;
28 :
29 : /**
30 : * Generates SQL DDL from pg_attribute query results.
31 : * Builds a hierarchy of schema -> table -> columns and generates CREATE TABLE statements.
32 : */
33 1 : export function generateCatalogSQLFromPgAttribute(result: PgAttributeColumnsTable, databaseName: string): string {
34 : // Build hierarchy: schema -> table -> columns
35 1 : const schemas = new Map<string, Map<string, ColumnMetadata[]>>();
36 :
37 1 : for (const batch of result.batches) {
38 1 : const colTableSchema = batch.getChild("table_schema")!;
39 1 : const colTableName = batch.getChild("table_name")!;
40 1 : const colColumnName = batch.getChild("column_name")!;
41 1 : const colOrdinalPos = batch.getChild("ordinal_position")!;
42 1 : const colDataType = batch.getChild("data_type")!;
43 :
44 : // Iterate over all rows in the batch
45 1 : for (let i = 0; i < batch.numRows; ++i) {
46 1 : const tableSchema = colTableSchema.at(i);
47 1 : const tableName = colTableName.at(i);
48 1 : const columnName = colColumnName.at(i);
49 1 : const ordinalPosition = colOrdinalPos.at(i);
50 1 : const dataType = colDataType.at(i);
51 :
52 1 : if (!tableSchema || !columnName || !tableName) {
53 1 : continue;
54 : }
55 :
56 : // Get or create schema entry
57 1 : let tableMap = schemas.get(tableSchema);
58 1 : if (!tableMap) {
59 1 : tableMap = new Map<string, ColumnMetadata[]>();
60 1 : schemas.set(tableSchema, tableMap);
61 1 : }
62 :
63 : // Get or create table entry
64 1 : let columns = tableMap.get(tableName);
65 1 : if (!columns) {
66 1 : columns = [];
67 1 : tableMap.set(tableName, columns);
68 1 : }
69 :
70 : // Add column metadata
71 1 : columns.push({
72 1 : name: columnName,
73 1 : ordinalPosition: ordinalPosition ?? columns.length,
74 1 : dataType: dataType ?? null,
75 1 : });
76 1 : }
77 1 : }
78 :
79 : // Generate SQL for all schemas
80 1 : const sqlStatements: string[] = [];
81 1 : for (const [schemaName, schemaTables] of schemas) {
82 1 : const sql = generateSchemaSQL(databaseName, schemaName, schemaTables);
83 1 : if (sql.length > 0) {
84 1 : sqlStatements.push(sql);
85 : }
86 1 : }
87 :
88 1 : return sqlStatements.join('\n\n');
89 1 : }
90 :
91 1 : export async function queryPgAttribute(
92 1 : connectionId: string,
93 1 : connectionDispatch: DynamicConnectionDispatch,
94 1 : updateId: number,
95 1 : databaseName: string,
96 1 : schemaNames: string[],
97 1 : executor: QueryExecutor,
98 1 : queryDatabaseName: string | null = null,
99 1 : abortSignal?: AbortSignal,
100 1 : ): Promise<PgAttributeColumnsTable | null> {
101 1 : const pgCatalog = queryDatabaseName == null
102 1 : ? 'pg_catalog'
103 1 : : `${quoteIdentifier(queryDatabaseName)}.${quoteIdentifier('pg_catalog')}`;
104 1 : const schemaFilter = schemaNames.length > 0
105 1 : ? `AND n.nspname IN (${schemaNames.map(name => `'${name.replace(/'/g, "''")}'`).join(', ')})`
106 1 : : '';
107 : const query = `
108 : SELECT
109 : n.nspname AS table_schema,
110 : c.relname AS table_name,
111 : a.attname AS column_name,
112 : a.attnum AS ordinal_position,
113 : t.typname AS data_type,
114 : CASE
115 : WHEN a.attnotnull THEN 'NO'
116 : ELSE 'YES'
117 : END AS is_nullable,
118 : CASE
119 : WHEN a.atttypid = ANY (ARRAY[21, 23, 20]) THEN a.atttypmod - 4
120 : WHEN a.atttypid = 1700 THEN ((a.atttypmod - 4) >> 16) & 65535
121 : ELSE NULL
122 : END AS numeric_precision,
123 : CASE
124 : WHEN a.atttypid = 1700 THEN (a.atttypmod - 4) & 65535
125 : ELSE NULL
126 : END AS numeric_scale
127 1 : FROM ${pgCatalog}.pg_class c
128 1 : JOIN ${pgCatalog}.pg_namespace n ON n.oid = c.relnamespace
129 1 : JOIN ${pgCatalog}.pg_attribute a ON a.attrelid = c.oid
130 1 : JOIN ${pgCatalog}.pg_type t ON t.oid = a.atttypid
131 1 : LEFT JOIN ${pgCatalog}.pg_attrdef d ON d.adrelid = c.oid AND d.adnum = a.attnum
132 : WHERE c.relkind IN ('r', 'v', 'm', 'f', 'p')
133 : AND a.attnum > 0
134 : AND NOT a.attisdropped
135 1 : ${schemaFilter}
136 : ORDER BY n.nspname, c.relname, a.attnum
137 : `;
138 :
139 : const args: QueryExecutionArgs = {
140 1 : query: query,
141 1 : abortSignal,
142 1 : readTimeoutMs: CATALOG_QUERY_READ_TIMEOUT_MS,
143 1 : throwOnError: true,
144 1 : metadata: {
145 1 : queryType: QueryType.CATALOG_QUERY_PG_ATTRIBUTE,
146 1 : title: "Query Postgres Schema Columns",
147 1 : description: null,
148 1 : issuer: "Catalog Update",
149 1 : userProvided: false
150 1 : }
151 : };
152 1 : const [queryId, queryExecution] = executor(connectionId, args);
153 1 : connectionDispatch(connectionId, {
154 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
155 1 : value: [updateId, queryId]
156 1 : });
157 :
158 1 : const queryResult = await queryExecution as PgAttributeColumnsTable;
159 1 : return queryResult;
160 1 : }
161 :
162 1 : export async function updatePgSchemaScript(
163 1 : logger: LoggerLike,
164 1 : connectionId: string,
165 1 : connectionDispatch: DynamicConnectionDispatch,
166 1 : updateId: number,
167 1 : databaseName: string,
168 1 : schemaNames: string[],
169 1 : executor: QueryExecutor,
170 1 : catalog: dashql.DashQLCatalog,
171 1 : catalogRelationScript: dashql.DashQLScript,
172 1 : options: CatalogBatchOptions = {},
173 1 : ): Promise<void> {
174 1 : const validate = beginCatalogBatch(catalog, options);
175 1 : const queryResult = await queryPgAttribute(connectionId, connectionDispatch, updateId, databaseName, schemaNames, executor, null, options.abortSignal);
176 1 : validate();
177 1 : if (queryResult == null || queryResult.numRows === 0) {
178 1 : throw new Error('pg_attribute returned no catalog relations');
179 : }
180 :
181 1 : const header = generateCatalogScriptHeader(CatalogSource.PgClass);
182 1 : const catalogSQL = generateCatalogSQLFromPgAttribute(queryResult, databaseName);
183 1 : if (catalogSQL.length === 0) {
184 0 : throw new Error('pg_attribute returned no usable catalog relations');
185 : }
186 :
187 1 : connectionDispatch(connectionId, {
188 1 : type: CATALOG_UPDATE_SCHEMA_SCRIPT,
189 1 : value: [updateId]
190 1 : });
191 :
192 1 : const [relations] = await publishCatalogBatch(catalogRelationScript.ptr.api, catalog, [{
193 1 : script: catalogRelationScript, text: `${header}${catalogSQL}`,
194 1 : }], options, validate);
195 1 : logger.info("Collected tables from pg_attribute", {
196 1 : "updateId": updateId.toString(),
197 1 : "tables": relations.tables.toString(),
198 1 : }, LOG_CTX);
199 1 : }
200 :
201 1 : export async function updatePgCatalog(
202 1 : logger: LoggerLike,
203 1 : connectionId: string,
204 1 : connectionDispatch: DynamicConnectionDispatch,
205 1 : updateId: number,
206 1 : databaseName: string,
207 1 : schemaNames: string[],
208 1 : executor: QueryExecutor,
209 1 : catalog: dashql.DashQLCatalog,
210 1 : dql: dashql.DashQL,
211 1 : catalogRelationScript: dashql.DashQLScript,
212 1 : catalogFunctionScript: dashql.DashQLScript,
213 1 : options: CatalogBatchOptions = {},
214 1 : ): Promise<void> {
215 1 : const validate = beginCatalogBatch(catalog, options);
216 1 : const [relationsResult, functionsResult] = await Promise.all([
217 1 : queryPgAttribute(connectionId, connectionDispatch, updateId, databaseName, schemaNames, executor, null, options.abortSignal),
218 1 : queryPgProc(connectionId, connectionDispatch, updateId, executor, options.abortSignal),
219 1 : ]);
220 1 : validate();
221 1 : if (relationsResult == null || relationsResult.numRows === 0) {
222 0 : throw new Error('pg_attribute returned no catalog relations');
223 : }
224 1 : const relationsSQL = generateCatalogSQLFromPgAttribute(relationsResult, databaseName);
225 1 : if (!relationsSQL) throw new Error('pg_attribute returned no usable catalog relations');
226 1 : const entries: CatalogBatchEntry[] = [{
227 1 : script: catalogRelationScript,
228 1 : text: `${generateCatalogScriptHeader(CatalogSource.PgClass)}${relationsSQL}`,
229 1 : }];
230 1 : if (functionsResult == null) throw new Error('pg_proc returned no function metadata result');
231 1 : entries.push({
232 1 : script: catalogFunctionScript,
233 1 : text: `${generateFunctionScriptHeader(CatalogSource.PgClass)}${generateCatalogSQLFromPgProc(functionsResult, databaseName, logger)}`,
234 1 : });
235 1 : connectionDispatch(connectionId, { type: CATALOG_UPDATE_SCHEMA_SCRIPT, value: [updateId] });
236 1 : const counts = await publishCatalogBatch(dql, catalog, entries, options, validate);
237 1 : logger.info('Collected PostgreSQL catalog', {
238 1 : updateId: updateId.toString(),
239 1 : tables: counts[0].tables.toString(),
240 1 : functions: counts[1].functions.toString(),
241 1 : }, LOG_CTX);
242 1 : }
|
