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_REGISTER_QUERY } from "../connection_state.js";
8 : import { QueryType } from "../../query/query_execution_state.js";
9 : import { catalogTablesOfSchema, packCatalogDescriptor, packCatalogTable, packRelationName, quoteIdentifier, quoteLiteral, relationFilter, type CatalogTable, type ColumnMetadata, type RelationName } from './catalog_sql_generator.js';
10 : import { beginCatalogBatch, 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 : data_type: arrow.Utf8;
20 : }>;
21 :
22 : /**
23 : * The tables of information_schema query results.
24 : * Builds a hierarchy of catalog -> schema -> table -> columns.
25 : */
26 1 : function catalogTablesFromInformationSchema(result: InformationSchemaColumnsTable): CatalogTable[] {
27 : // Build hierarchy: catalog -> schema -> table -> columns
28 1 : const catalogs = new Map<string, Map<string, Map<string, ColumnMetadata[]>>>();
29 :
30 1 : for (const batch of result.batches) {
31 1 : const colTableCatalog = batch.getChild("table_catalog")!;
32 1 : const colTableSchema = batch.getChild("table_schema")!;
33 1 : const colTableName = batch.getChild("table_name")!;
34 1 : const colColumnName = batch.getChild("column_name")!;
35 1 : const colOrdinalPos = batch.getChild("ordinal_position")!;
36 1 : const colDataType = batch.getChild("data_type")!;
37 :
38 : // Iterate over all rows in the batch
39 1 : for (let i = 0; i < batch.numRows; ++i) {
40 1 : const tableCatalog = colTableCatalog.at(i);
41 1 : const tableSchema = colTableSchema.at(i);
42 1 : const tableName = colTableName.at(i);
43 1 : const columnName = colColumnName.at(i);
44 1 : const ordinalPosition = colOrdinalPos.at(i);
45 1 : const dataType = colDataType.at(i);
46 :
47 1 : if (!tableCatalog || !tableSchema || !columnName || !tableName) {
48 0 : continue;
49 : }
50 :
51 : // Get or create catalog entry
52 1 : let schemaMap = catalogs.get(tableCatalog);
53 1 : if (!schemaMap) {
54 1 : schemaMap = new Map<string, Map<string, ColumnMetadata[]>>();
55 1 : catalogs.set(tableCatalog, schemaMap);
56 1 : }
57 :
58 : // Get or create schema entry
59 1 : let tableMap = schemaMap.get(tableSchema);
60 1 : if (!tableMap) {
61 1 : tableMap = new Map<string, ColumnMetadata[]>();
62 1 : schemaMap.set(tableSchema, tableMap);
63 1 : }
64 :
65 : // Get or create table entry
66 1 : let columns = tableMap.get(tableName);
67 1 : if (!columns) {
68 1 : columns = [];
69 1 : tableMap.set(tableName, columns);
70 1 : }
71 :
72 : // Add column metadata
73 1 : columns.push({
74 1 : name: columnName,
75 1 : ordinalPosition: ordinalPosition != null ? Number(ordinalPosition) : columns.length,
76 1 : dataType: dataType ?? null,
77 1 : });
78 1 : }
79 1 : }
80 :
81 1 : return [...catalogs].flatMap(([catalogName, schemas]) =>
82 1 : [...schemas].flatMap(([schemaName, tables]) => catalogTablesOfSchema(catalogName, schemaName, tables)));
83 1 : }
84 :
85 : /// A catalog and the schemas of it whose columns are read, no schemas for all of them
86 : export interface InformationSchemaScope {
87 : /// The catalog name
88 : catalog: string;
89 : /// The schema names, empty for all schemas of the catalog
90 : schemas: string[];
91 : /// The relations, all relations of the schemas if absent
92 : relations?: RelationName[];
93 : }
94 :
95 : /// The query of the columns of the scopes.
96 : /// Every catalog has an information_schema of its own, the query reads the one of every catalog.
97 1 : export function informationSchemaColumnsQuery(scopes: InformationSchemaScope[]): string {
98 1 : const condition = (scope: InformationSchemaScope) => {
99 1 : if (scope.relations != null) return `\nWHERE table_schema <> 'information_schema' AND ${relationFilter('table_schema', 'table_name', scope.relations)}`;
100 1 : return scope.schemas.length > 0 ? `\nWHERE table_schema IN (${scope.schemas.map(quoteLiteral).join(', ')})` : '';
101 1 : };
102 1 : const select = (source: string, scope: InformationSchemaScope) => `SELECT
103 : table_catalog,
104 : table_schema,
105 : table_name,
106 : column_name,
107 : ordinal_position,
108 : data_type
109 1 : FROM ${source}${condition(scope)}`;
110 1 : return scopes
111 1 : .map(scope => select(`${quoteIdentifier(scope.catalog)}.information_schema.columns`, scope))
112 1 : .join('\nUNION ALL\n');
113 1 : }
114 :
115 1 : export async function queryInformationSchema(connectionId: string, connectionDispatch: DynamicConnectionDispatch, updateId: number, scopes: InformationSchemaScope[], executor: QueryExecutor, abortSignal?: AbortSignal): Promise<InformationSchemaColumnsTable | null> {
116 1 : const query = informationSchemaColumnsQuery(scopes);
117 :
118 : const args: QueryExecutionArgs = {
119 1 : query: query,
120 1 : abortSignal,
121 1 : throwOnError: true,
122 1 : metadata: {
123 1 : queryType: QueryType.CATALOG_QUERY_INFORMATION_SCHEMA,
124 1 : title: "Information Schema",
125 1 : description: null,
126 1 : issuer: "Catalog Update",
127 1 : userProvided: false
128 1 : }
129 : };
130 1 : const [queryId, queryExecution] = executor(connectionId, args);
131 1 : connectionDispatch(connectionId, {
132 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
133 1 : value: [updateId, queryId]
134 1 : });
135 :
136 1 : const queryResult = await queryExecution as InformationSchemaColumnsTable;
137 1 : return queryResult;
138 1 : }
139 :
140 : /// The scopes of relations: their catalogs
141 1 : export function relationScopes(relations: RelationName[]): InformationSchemaScope[] {
142 1 : const catalogs = [...new Set(relations.map(relation => relation.databaseName))];
143 1 : return catalogs.map(catalog => ({ catalog, schemas: [], relations: relations.filter(relation => relation.databaseName === catalog) }));
144 1 : }
145 :
146 : /// Ask the information_schema about relations and update them in the catalog
147 1 : export async function updateInformationSchemaRelations(
148 1 : connectionId: string,
149 1 : connectionDispatch: DynamicConnectionDispatch,
150 1 : updateId: number,
151 1 : relations: RelationName[],
152 1 : executor: QueryExecutor,
153 1 : catalog: dashql.DashQLCatalog,
154 1 : relationEntryId: number,
155 1 : options: CatalogBatchOptions = {},
156 1 : ): Promise<void> {
157 1 : const validate = beginCatalogBatch(catalog, options);
158 1 : const queryResult = await queryInformationSchema(connectionId, connectionDispatch, updateId, relationScopes(relations), executor, options.abortSignal);
159 1 : validate();
160 1 : const tables = queryResult == null ? [] : catalogTablesFromInformationSchema(queryResult);
161 1 : catalog.replaceTables(relationEntryId, dashql.buffers.catalog.CatalogLayer.CONNECTION, {
162 1 : relations: relations.map(packRelationName),
163 1 : tables: tables.map(packCatalogTable),
164 1 : });
165 1 : }
166 :
167 1 : export async function updateInformationSchemaCatalog(
168 1 : connectionId: string,
169 1 : connectionDispatch: DynamicConnectionDispatch,
170 1 : updateId: number,
171 1 : scopes: InformationSchemaScope[],
172 1 : executor: QueryExecutor,
173 1 : catalog: dashql.DashQLCatalog,
174 1 : relationEntryId: number,
175 1 : options: CatalogBatchOptions = {},
176 1 : ): Promise<void> {
177 : // Every information_schema belongs to a catalog, without one there is nothing to read
178 1 : if (scopes.length === 0) {
179 1 : throw new Error('no catalogs to read the information_schema of');
180 : }
181 1 : const validate = beginCatalogBatch(catalog, options);
182 : // Query the information schema. If the query errors it throws and propagates
183 : // to the caller so we never overwrite the existing catalog with partial data.
184 1 : const queryResult = await queryInformationSchema(connectionId, connectionDispatch, updateId, scopes, executor, options.abortSignal);
185 1 : validate();
186 1 : if (queryResult == null || queryResult.numRows === 0) {
187 1 : throw new Error('information_schema returned no catalog relations');
188 : }
189 : // An empty result after row-level filtering doesn't clobber a restored catalog
190 1 : const tables = catalogTablesFromInformationSchema(queryResult);
191 1 : if (tables.length === 0) {
192 0 : throw new Error('information_schema returned no usable catalog relations');
193 : }
194 1 : catalog.replaceDescriptor(relationEntryId, dashql.buffers.catalog.CatalogLayer.CONNECTION, packCatalogDescriptor(tables));
195 1 : }
|
