Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import * as dashql from '@ankoh/dashql-core';
3 :
4 : import { CATALOG_UPDATE_REGISTER_QUERY } from '../connection_state.js';
5 : import { DynamicConnectionDispatch } from '../connection_registry.js';
6 : import { catalogTablesOfSchema, DEFAULT_DATABASE_NAME, packCatalogDescriptor, packCatalogTable, packRelationName, quoteIdentifier, relationFilter, type CatalogTable, type ColumnMetadata, type RelationName } from '../catalog/catalog_sql_generator.js';
7 : import { CATALOG_QUERY_READ_TIMEOUT_MS, queryPgAttribute, catalogTablesFromPgAttribute } from '../catalog/catalog_query_pg_attribute.js';
8 : import { QueryExecutionArgs } from '../query_execution_args.js';
9 : import { QueryExecutor } from '../query_executor.js';
10 : import { QueryType } from '../../query/query_execution_state.js';
11 : import type { LoggerLike } from '../../platform/logger/logger.js';
12 : import type { AttachedDatabase } from './hyperdb_grpc_client.js';
13 : import { qualifyPrefetchedHyperFunctions } from '../catalog/prefetched_hyper_functions.js';
14 : import { beginCatalogBatch, publishCatalogBatch, type CatalogBatchOptions } from '../catalog/catalog_batch.js';
15 : import { generateFunctionScriptHeaderForSource } from '../catalog/catalog_function_sql_generator.js';
16 : import { generateCatalogSQLFromPgProc, queryPgProc } from '../catalog/catalog_query_pg_proc.js';
17 :
18 1 : const LOG_CTX = 'hyper_catalog';
19 :
20 : export type HyperCloudCatalogTable = arrow.Table<{
21 : schema_name: arrow.Utf8;
22 : table_name: arrow.Utf8;
23 : column_name: arrow.Utf8;
24 : ordinal_position: arrow.Int32;
25 : data_type: arrow.Utf8;
26 : }>;
27 :
28 : export interface HyperCatalogUpdateFailure {
29 : database: string;
30 : error: Error;
31 : }
32 :
33 : export interface HyperCatalogUpdateResult {
34 : updatedDatabases: string[];
35 : failures: HyperCatalogUpdateFailure[];
36 : }
37 :
38 : interface CatalogTarget {
39 : key: string;
40 : databaseName: string;
41 : queryDatabaseName: string | null;
42 : path: string;
43 : }
44 :
45 1 : function toError(value: unknown): Error {
46 1 : return value instanceof Error ? value : new Error(String(value));
47 1 : }
48 :
49 1 : function summarizeFailures(message: string, failures: HyperCatalogUpdateFailure[]): Error {
50 1 : return new Error(`${message}: ${failures.map(failure => `${failure.database}: ${failure.error.message}`).join('; ')}`);
51 1 : }
52 :
53 1 : function buildCatalogTargets(attachedDatabases: AttachedDatabase[]): CatalogTarget[] {
54 1 : if (attachedDatabases.length === 0) {
55 1 : return [{ key: DEFAULT_DATABASE_NAME, databaseName: DEFAULT_DATABASE_NAME, queryDatabaseName: null, path: '' }];
56 : }
57 1 : return attachedDatabases.map(database => ({
58 1 : key: database.alias ?? '',
59 1 : databaseName: database.alias ?? '',
60 1 : queryDatabaseName: database.alias || null,
61 1 : path: database.path,
62 1 : }));
63 1 : }
64 :
65 1 : export function buildHyperCloudCatalogQuery(databaseAlias: string, relations: RelationName[] | null = null): string {
66 1 : const database = quoteIdentifier(databaseAlias);
67 1 : const table = (name: string) => `${database}.${quoteIdentifier('_hyper_catalog')}.${quoteIdentifier(name)}`;
68 1 : return `
69 : SELECT
70 : s.schema_name_display AS schema_name,
71 : o.object_name_display AS table_name,
72 : c.column_name_display AS column_name,
73 : c.ordinal_position AS ordinal_position,
74 : CAST(c.type_descriptor AS TEXT) AS data_type
75 1 : FROM ${table('databases')} d
76 1 : JOIN ${table('schemas')} s
77 : ON s.account_id = d.account_id
78 : AND s.database_id = d.database_id
79 1 : JOIN ${table('objects')} o
80 : ON o.account_id = s.account_id
81 : AND o.schema_id = s.schema_id
82 1 : JOIN ${table('columns')} c
83 : ON c.account_id = o.account_id
84 : AND c.object_id = o.object_id
85 : WHERE o.object_type IN ('ICEBERG_TABLE', 'VIEW')
86 1 : ${relations == null ? '' : `AND ${relationFilter('s.schema_name_display', 'o.object_name_display', relations)}`}
87 : ORDER BY
88 : s.schema_name_display,
89 : o.object_name_display,
90 : c.ordinal_position
91 : `;
92 1 : }
93 :
94 1 : export function catalogTablesFromHyperCloud(result: HyperCloudCatalogTable, connectionName: string, logger?: LoggerLike): CatalogTable[] {
95 1 : const schemas = new Map<string, Map<string, ColumnMetadata[]>>();
96 1 : const seenColumns = new Set<string>();
97 1 : let duplicateColumns = 0;
98 1 : let firstDuplicate: Record<string, string> | null = null;
99 1 : for (const batch of result.batches) {
100 1 : const schemaNames = batch.getChild('schema_name')!;
101 1 : const tableNames = batch.getChild('table_name')!;
102 1 : const columnNames = batch.getChild('column_name')!;
103 1 : const ordinalPositions = batch.getChild('ordinal_position')!;
104 1 : const dataTypes = batch.getChild('data_type')!;
105 :
106 1 : for (let i = 0; i < batch.numRows; ++i) {
107 1 : const schemaName = schemaNames.at(i);
108 1 : const tableName = tableNames.at(i);
109 1 : const columnName = columnNames.at(i);
110 1 : if (!schemaName || !tableName || !columnName) continue;
111 :
112 1 : let tables = schemas.get(schemaName);
113 1 : if (!tables) schemas.set(schemaName, tables = new Map());
114 1 : let columns = tables.get(tableName);
115 1 : if (!columns) tables.set(tableName, columns = []);
116 : // The catalog can list a column more than once, a table declares it once
117 1 : const columnKey = JSON.stringify([schemaName, tableName, columnName]);
118 1 : if (seenColumns.has(columnKey)) {
119 1 : firstDuplicate ??= {
120 1 : schema: schemaName, table: tableName, column: columnName,
121 1 : ordinalPosition: String(ordinalPositions.at(i)), dataType: String(dataTypes.at(i)),
122 1 : };
123 1 : ++duplicateColumns;
124 1 : continue;
125 1 : }
126 1 : seenColumns.add(columnKey);
127 1 : columns.push({
128 1 : name: columnName,
129 1 : ordinalPosition: ordinalPositions.at(i) ?? columns.length,
130 1 : dataType: dataTypes.at(i) ?? null,
131 1 : });
132 1 : }
133 1 : }
134 1 : if (firstDuplicate != null) {
135 1 : logger?.warn('Skipping duplicate Hyper Cloud catalog columns', {
136 1 : database: connectionName, duplicates: duplicateColumns.toString(), ...firstDuplicate,
137 1 : }, LOG_CTX);
138 : }
139 :
140 1 : return [...schemas.keys()].sort().flatMap(schemaName => catalogTablesOfSchema(connectionName, schemaName, schemas.get(schemaName)!));
141 1 : }
142 :
143 1 : async function queryHyperCloudCatalog(
144 1 : connectionId: string,
145 1 : connectionDispatch: DynamicConnectionDispatch,
146 1 : updateId: number,
147 1 : databaseAlias: string,
148 1 : executor: QueryExecutor,
149 1 : abortSignal: AbortSignal,
150 1 : relations: RelationName[] | null,
151 1 : ): Promise<HyperCloudCatalogTable> {
152 : const args: QueryExecutionArgs = {
153 1 : query: buildHyperCloudCatalogQuery(databaseAlias, relations),
154 1 : abortSignal,
155 1 : readTimeoutMs: CATALOG_QUERY_READ_TIMEOUT_MS,
156 1 : throwOnError: true,
157 1 : metadata: {
158 1 : queryType: QueryType.CATALOG_QUERY_HYPER_CLOUD,
159 1 : title: `Query Hyper Cloud Catalog (${databaseAlias})`,
160 1 : description: null,
161 1 : issuer: 'Catalog Update',
162 1 : userProvided: false,
163 1 : },
164 : };
165 1 : const [queryId, queryExecution] = executor(connectionId, args);
166 1 : connectionDispatch(connectionId, {
167 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
168 1 : value: [updateId, queryId],
169 1 : });
170 1 : const result = await queryExecution as HyperCloudCatalogTable | null;
171 1 : if (result == null || (relations == null && result.numRows === 0)) {
172 0 : throw new Error(`Hyper Cloud catalog for ${databaseAlias} returned no relations`);
173 : }
174 1 : return result;
175 1 : }
176 :
177 1 : export async function updateHyperCatalog(
178 1 : logger: LoggerLike,
179 1 : connectionId: string,
180 1 : connectionDispatch: DynamicConnectionDispatch,
181 1 : updateId: number,
182 1 : attachedDatabases: AttachedDatabase[],
183 1 : executor: QueryExecutor,
184 1 : catalog: dashql.DashQLCatalog,
185 1 : dql: dashql.DashQL,
186 1 : relationEntryId: number,
187 1 : catalogFunctionScript: dashql.DashQLScript,
188 1 : resolveLiveFunctionCatalog: boolean,
189 1 : abortSignal: AbortSignal,
190 1 : options: CatalogBatchOptions = {},
191 : /// The relations to ask about, all relations and the functions if absent
192 1 : relations: RelationName[] | null = null,
193 1 : ): Promise<HyperCatalogUpdateResult> {
194 1 : options = { ...options, abortSignal };
195 1 : const validate = beginCatalogBatch(catalog, options);
196 1 : const targets = buildCatalogTargets(attachedDatabases);
197 1 : const aliasCounts = new Map<string, number>();
198 1 : for (const target of targets) aliasCounts.set(target.key, (aliasCounts.get(target.key) ?? 0) + 1);
199 :
200 1 : const relationsOf = (target: CatalogTarget) => relations?.filter(relation => relation.databaseName === target.databaseName) ?? null;
201 1 : const queriedTargets = relations == null ? targets : targets.filter(target => relationsOf(target)!.length > 0);
202 1 : const results = await Promise.all(queriedTargets.map(async target => {
203 1 : try {
204 1 : abortSignal.throwIfAborted();
205 1 : if (attachedDatabases.length > 0 && !target.key.trim()) {
206 1 : throw new Error('Attached databases require a non-empty alias for catalog refresh');
207 : }
208 1 : if ((aliasCounts.get(target.key) ?? 0) > 1) {
209 1 : throw new Error(`Attached database alias ${JSON.stringify(target.key)} is not unique`);
210 : }
211 :
212 1 : let tables: CatalogTable[];
213 1 : if (target.path.startsWith('hyper.cloud')) {
214 1 : const result = await queryHyperCloudCatalog(
215 1 : connectionId,
216 1 : connectionDispatch,
217 1 : updateId,
218 1 : target.databaseName,
219 1 : executor,
220 1 : abortSignal,
221 1 : relationsOf(target),
222 1 : );
223 1 : tables = catalogTablesFromHyperCloud(result, target.databaseName, logger);
224 1 : } else {
225 1 : const result = await queryPgAttribute(
226 1 : connectionId,
227 1 : connectionDispatch,
228 1 : updateId,
229 1 : executor,
230 1 : target.queryDatabaseName,
231 1 : abortSignal,
232 1 : relationsOf(target),
233 1 : );
234 1 : if (result == null || (relations == null && result.numRows === 0)) {
235 1 : throw new Error(`pg_attribute for ${target.databaseName} returned no relations`);
236 : }
237 1 : tables = catalogTablesFromPgAttribute(result, target.databaseName);
238 1 : }
239 1 : if (relations == null && tables.length === 0) throw new Error(`Catalog for ${target.databaseName} returned no usable relations`);
240 1 : return { target, tables, error: null };
241 1 : } catch (error) {
242 1 : return { target, tables: null, error: toError(error) };
243 1 : }
244 1 : }));
245 :
246 1 : validate();
247 1 : const successful = results.filter(result => result.tables != null);
248 1 : const failures = results
249 1 : .filter(result => result.error != null)
250 1 : .map(result => ({ database: result.target.key || result.target.path, error: result.error! }));
251 1 : if (successful.length === 0 && (relations == null || failures.length > 0)) {
252 1 : throw summarizeFailures('Failed to refresh every attached database', failures);
253 : }
254 1 : const updatedDatabases = successful.map(result => result.target.key);
255 1 : const tables = successful.flatMap(result => result.tables!);
256 1 : const layer = dashql.buffers.catalog.CatalogLayer.CONNECTION;
257 :
258 : // The answer about relations updates them, the rest of the catalog stays
259 1 : if (relations != null) {
260 1 : const answered = successful.flatMap(result => relationsOf(result.target)!);
261 1 : if (answered.length > 0) {
262 1 : catalog.replaceTables(relationEntryId, layer, { relations: answered.map(packRelationName), tables: tables.map(packCatalogTable) });
263 : }
264 1 : return { updatedDatabases, failures };
265 1 : }
266 :
267 1 : let functionSQL: string;
268 1 : if (resolveLiveFunctionCatalog) {
269 1 : const functionResult = await queryPgProc(
270 1 : connectionId,
271 1 : connectionDispatch,
272 1 : updateId,
273 1 : executor,
274 1 : abortSignal,
275 1 : );
276 1 : validate();
277 1 : if (functionResult == null) {
278 1 : throw new Error('pg_proc returned no Hyper function metadata result');
279 : }
280 1 : const functions = generateCatalogSQLFromPgProc(functionResult, DEFAULT_DATABASE_NAME, logger);
281 1 : functionSQL = `${generateFunctionScriptHeaderForSource('HyperDB WASM pg_proc')}${functions}`;
282 1 : } else {
283 1 : functionSQL = qualifyPrefetchedHyperFunctions(DEFAULT_DATABASE_NAME);
284 : }
285 :
286 1 : validate();
287 1 : const [functions] = await publishCatalogBatch(dql, catalog, [
288 1 : { script: catalogFunctionScript, text: functionSQL, requireFunctions: !resolveLiveFunctionCatalog },
289 1 : ], options, validate);
290 1 : validate();
291 : // A database that failed keeps the relations it had
292 1 : if (failures.length === 0) {
293 1 : catalog.replaceDescriptor(relationEntryId, layer, packCatalogDescriptor(tables));
294 : } else {
295 1 : catalog.replaceTables(relationEntryId, layer, { databases: successful.map(result => result.target.databaseName), tables: tables.map(packCatalogTable) });
296 : }
297 1 : logger.info('Updated Hyper catalog relations', {
298 1 : updateId: updateId.toString(),
299 1 : databasesUpdated: successful.length.toString(),
300 1 : databasesFailed: failures.length.toString(),
301 1 : tables: tables.length.toString(),
302 1 : functions: functions.functions.toString(),
303 1 : }, LOG_CTX);
304 1 : return { updatedDatabases, failures };
305 1 : }
|
