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, type ColumnMetadata } from './catalog_sql_generator.js';
10 : import { beginCatalogBatch, publishCatalogBatch, type CatalogBatchOptions } from './catalog_batch.js';
11 :
12 : export type InformationSchemaColumnsTable = arrow.Table<{
13 : table_catalog: arrow.Utf8;
14 : table_schema: arrow.Utf8;
15 : table_name: arrow.Utf8;
16 : column_name: arrow.Utf8;
17 : /// Trino types the ordinal position as bigint, its values are then BigInt
18 : ordinal_position: arrow.Int32 | arrow.Int64;
19 : is_nullable: arrow.Utf8; // 'YES' or 'NO'
20 : data_type: arrow.Utf8;
21 : }>;
22 :
23 : /**
24 : * Generates SQL DDL from information_schema query results.
25 : * Builds a hierarchy of catalog -> schema -> table -> columns and generates CREATE TABLE statements.
26 : */
27 1 : function generateCatalogSQLFromInformationSchema(result: InformationSchemaColumnsTable): string {
28 : // Build hierarchy: catalog -> schema -> table -> columns
29 1 : const catalogs = new Map<string, Map<string, Map<string, ColumnMetadata[]>>>();
30 :
31 1 : for (const batch of result.batches) {
32 1 : const colTableCatalog = batch.getChild("table_catalog")!;
33 1 : const colTableSchema = batch.getChild("table_schema")!;
34 1 : const colTableName = batch.getChild("table_name")!;
35 1 : const colColumnName = batch.getChild("column_name")!;
36 1 : const colOrdinalPos = batch.getChild("ordinal_position")!;
37 1 : const colDataType = batch.getChild("data_type")!;
38 :
39 : // Iterate over all rows in the batch
40 1 : for (let i = 0; i < batch.numRows; ++i) {
41 1 : const tableCatalog = colTableCatalog.at(i);
42 1 : const tableSchema = colTableSchema.at(i);
43 1 : const tableName = colTableName.at(i);
44 1 : const columnName = colColumnName.at(i);
45 1 : const ordinalPosition = colOrdinalPos.at(i);
46 1 : const dataType = colDataType.at(i);
47 :
48 1 : if (!tableCatalog || !tableSchema || !columnName || !tableName) {
49 0 : continue;
50 : }
51 :
52 : // Get or create catalog entry
53 1 : let schemaMap = catalogs.get(tableCatalog);
54 1 : if (!schemaMap) {
55 1 : schemaMap = new Map<string, Map<string, ColumnMetadata[]>>();
56 1 : catalogs.set(tableCatalog, schemaMap);
57 1 : }
58 :
59 : // Get or create schema entry
60 1 : let tableMap = schemaMap.get(tableSchema);
61 1 : if (!tableMap) {
62 1 : tableMap = new Map<string, ColumnMetadata[]>();
63 1 : schemaMap.set(tableSchema, tableMap);
64 1 : }
65 :
66 : // Get or create table entry
67 1 : let columns = tableMap.get(tableName);
68 1 : if (!columns) {
69 1 : columns = [];
70 1 : tableMap.set(tableName, columns);
71 1 : }
72 :
73 : // Add column metadata
74 1 : columns.push({
75 1 : name: columnName,
76 1 : ordinalPosition: ordinalPosition != null ? Number(ordinalPosition) : columns.length,
77 1 : dataType: dataType ?? null,
78 1 : });
79 1 : }
80 1 : }
81 :
82 : // Generate SQL for all schemas
83 1 : const sqlStatements: string[] = [];
84 1 : for (const [catalogName, catalogSchemas] of catalogs) {
85 1 : for (const [schemaName, schemaTables] of catalogSchemas) {
86 1 : const sql = generateSchemaSQL(catalogName, schemaName, schemaTables);
87 1 : if (sql.length > 0) {
88 1 : sqlStatements.push(sql);
89 : }
90 1 : }
91 : }
92 :
93 1 : return sqlStatements.join('\n\n');
94 1 : }
95 :
96 : /// A catalog and the schemas of it whose columns are read, no schemas for all of them
97 : export interface InformationSchemaScope {
98 : /// The catalog name
99 : catalog: string;
100 : /// The schema names, empty for all schemas of the catalog
101 : schemas: string[];
102 : }
103 :
104 : /// Quote an identifier
105 1 : function quoteIdentifier(name: string): string {
106 1 : return `"${name.replace(/"/g, '""')}"`;
107 1 : }
108 :
109 : /// Quote a string literal
110 1 : function quoteLiteral(value: string): string {
111 1 : return `'${value.replace(/'/g, "''")}'`;
112 1 : }
113 :
114 : /// The query of the columns of the scopes.
115 : /// Every catalog has an information_schema of its own, the query reads the one of every catalog.
116 1 : export function informationSchemaColumnsQuery(scopes: InformationSchemaScope[]): string {
117 1 : const select = (source: string, schemas: string[]) => `SELECT
118 : table_catalog,
119 : table_schema,
120 : table_name,
121 : column_name,
122 : ordinal_position,
123 : is_nullable,
124 : data_type
125 1 : FROM ${source}${schemas.length > 0 ? `
126 1 : WHERE table_schema IN (${schemas.map(quoteLiteral).join(', ')})` : ''}`;
127 1 : return scopes
128 1 : .map(scope => select(`${quoteIdentifier(scope.catalog)}.information_schema.columns`, scope.schemas))
129 1 : .join('\nUNION ALL\n');
130 1 : }
131 :
132 1 : export async function queryInformationSchema(connectionId: string, connectionDispatch: DynamicConnectionDispatch, updateId: number, scopes: InformationSchemaScope[], executor: QueryExecutor, abortSignal?: AbortSignal): Promise<InformationSchemaColumnsTable | null> {
133 1 : const query = informationSchemaColumnsQuery(scopes);
134 :
135 : const args: QueryExecutionArgs = {
136 1 : query: query,
137 1 : abortSignal,
138 1 : throwOnError: true,
139 1 : metadata: {
140 1 : queryType: QueryType.CATALOG_QUERY_INFORMATION_SCHEMA,
141 1 : title: "Information Schema",
142 1 : description: null,
143 1 : issuer: "Catalog Update",
144 1 : userProvided: false
145 1 : }
146 : };
147 1 : const [queryId, queryExecution] = executor(connectionId, args);
148 1 : connectionDispatch(connectionId, {
149 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
150 1 : value: [updateId, queryId]
151 1 : });
152 :
153 1 : const queryResult = await queryExecution as InformationSchemaColumnsTable;
154 1 : return queryResult;
155 1 : }
156 :
157 1 : export async function updateInformationSchemaCatalog(
158 1 : connectionId: string,
159 1 : connectionDispatch: DynamicConnectionDispatch,
160 1 : updateId: number,
161 1 : scopes: InformationSchemaScope[],
162 1 : executor: QueryExecutor,
163 1 : catalog: dashql.DashQLCatalog,
164 1 : dql: dashql.DashQL,
165 1 : catalogRelationScript: dashql.DashQLScript,
166 1 : _catalogFunctionScript: dashql.DashQLScript,
167 1 : options: CatalogBatchOptions = {},
168 1 : ): Promise<void> {
169 : // Every information_schema belongs to a catalog, without one there is nothing to read
170 1 : if (scopes.length === 0) {
171 1 : throw new Error('no catalogs to read the information_schema of');
172 : }
173 1 : const validate = beginCatalogBatch(catalog, options);
174 : // Query the information schema. If the query errors it throws and propagates
175 : // to the caller so we never overwrite the existing catalog script with partial data.
176 1 : const queryResult = await queryInformationSchema(connectionId, connectionDispatch, updateId, scopes, executor, options.abortSignal);
177 1 : validate();
178 1 : if (queryResult == null || queryResult.numRows === 0) {
179 1 : throw new Error('information_schema returned no catalog relations');
180 : }
181 :
182 : // Generate SQL from query results before touching the script so an empty
183 : // result (after row-level filtering) doesn't clobber a restored catalog.
184 1 : const header = generateCatalogScriptHeader(CatalogSource.InformationSchema);
185 1 : const catalogSQL = generateCatalogSQLFromInformationSchema(queryResult);
186 1 : if (catalogSQL.length === 0) {
187 0 : throw new Error('information_schema returned no usable catalog relations');
188 : }
189 :
190 : // Mark loading started
191 1 : connectionDispatch(connectionId, {
192 1 : type: CATALOG_UPDATE_SCHEMA_SCRIPT,
193 1 : value: [updateId]
194 1 : });
195 :
196 1 : await publishCatalogBatch(dql, catalog, [{
197 1 : script: catalogRelationScript, text: `${header}${catalogSQL}`,
198 1 : }], options, validate);
199 1 : }
|
