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