Line data Source code
1 : import * as arrow from "apache-arrow";
2 :
3 : import { QueryExecutor } from '../query_executor.js';
4 : import { QueryExecutionArgs } from '../query_execution_args.js';
5 : import { DynamicConnectionDispatch } from '../connection_registry.js';
6 : import { CATALOG_UPDATE_REGISTER_QUERY } from "../connection_state.js";
7 : import { QueryType } from '../../query/query_execution_state.js';
8 : import { type FunctionMetadata, generateFunctionsSQL } from './catalog_function_sql_generator.js';
9 : import { quoteIdentifier } from './catalog_sql_generator.js';
10 : import type { LoggerLike } from '../../platform/logger/logger.js';
11 :
12 1 : const CATALOG_QUERY_READ_TIMEOUT_MS = 60_000;
13 :
14 : export type PgProcTable = arrow.Table<{
15 : function_schema: arrow.Utf8;
16 : function_name: arrow.Utf8;
17 : argument_types: arrow.List<arrow.Utf8>;
18 : argument_names: arrow.List<arrow.Utf8>;
19 : argument_modes: arrow.List<arrow.Utf8>;
20 : return_type: arrow.Utf8;
21 : returns_set: arrow.Bool;
22 : function_kind: arrow.Utf8;
23 : }>;
24 :
25 1 : export function generateCatalogSQLFromPgProc(result: PgProcTable, databaseName: string | null | undefined, logger?: LoggerLike): string {
26 1 : const functions: FunctionMetadata[] = [];
27 :
28 1 : for (const batch of result.batches) {
29 1 : const colSchema = batch.getChild("function_schema")!;
30 1 : const colName = batch.getChild("function_name")!;
31 1 : const colTypes = batch.getChild('argument_types')!;
32 1 : const colNames = batch.getChild('argument_names')!;
33 1 : const colModes = batch.getChild('argument_modes')!;
34 1 : const colReturnsSet = batch.getChild('returns_set')!;
35 1 : const colReturn = batch.getChild("return_type")!;
36 1 : const colKind = batch.getChild("function_kind")!;
37 :
38 1 : for (let i = 0; i < batch.numRows; ++i) {
39 1 : const schemaName = colSchema.at(i);
40 1 : const functionName = colName.at(i);
41 1 : const returnType = colReturn.at(i);
42 1 : const kind = colKind.at(i);
43 :
44 1 : if (!functionName) {
45 1 : continue;
46 : }
47 1 : const types = Array.from(colTypes.at(i) ?? []);
48 1 : const names = Array.from(colNames.at(i) ?? []);
49 1 : const modes = Array.from(colModes.at(i) ?? []);
50 1 : if (!types.every(type => typeof type === 'string')
51 1 : || !names.every(name => name == null || typeof name === 'string')
52 1 : || !modes.every(mode => typeof mode === 'string')
53 1 : || (names.length !== 0 && names.length !== types.length)
54 1 : || (modes.length !== 0 && modes.length !== types.length)) {
55 1 : throw new Error('pg_proc returned invalid function argument metadata');
56 : }
57 : // DashQL declarations model scalar signatures only. Qualified custom types and
58 : // set/table results are not in its type grammar; do not let one poison the batch.
59 1 : const unsupportedType = (type: string) => type.replace(/"(?:[^"]|"")*"/g, '').includes('.');
60 1 : if (colReturnsSet.at(i) || modes.includes('t') || !returnType
61 1 : || unsupportedType(returnType) || types.some(unsupportedType)
62 1 : || modes.some(mode => !['i', 'o', 'b', 'v', 't'].includes(mode))) {
63 1 : logger?.warn('Skipping unsupported PostgreSQL function signature', {
64 1 : schema: schemaName, function: functionName, returnType,
65 1 : }, 'catalog_pg');
66 1 : continue;
67 1 : }
68 1 : const args = types.flatMap((type, index) => modes[index] === 'o' ? [] : [
69 1 : `${quoteIdentifier(names[index] || `arg${index + 1}`)} ${type}`,
70 1 : ]).join(', ');
71 1 : functions.push({
72 1 : schemaName,
73 1 : functionName,
74 1 : arguments: args,
75 1 : returnType,
76 1 : isAggregate: kind === 'a',
77 1 : });
78 1 : }
79 1 : }
80 :
81 1 : return generateFunctionsSQL(databaseName, functions);
82 1 : }
83 :
84 1 : export async function queryPgProc(
85 1 : connectionId: string,
86 1 : connectionDispatch: DynamicConnectionDispatch,
87 1 : updateId: number,
88 1 : executor: QueryExecutor,
89 1 : abortSignal?: AbortSignal,
90 1 : ): Promise<PgProcTable | null> {
91 : const query = `
92 : SELECT
93 : n.nspname AS function_schema,
94 : p.proname AS function_name,
95 : -- Hyper cannot cast oidvector to oid[]; format each source separately.
96 : CASE WHEN p.proallargtypes IS NOT NULL THEN ARRAY(
97 : SELECT pg_catalog.format_type(arg.type_oid, NULL)
98 : FROM unnest(p.proallargtypes) WITH ORDINALITY AS arg(type_oid, position)
99 : ORDER BY arg.position
100 : ) ELSE ARRAY(
101 : SELECT pg_catalog.format_type(arg.type_oid, NULL)
102 : FROM unnest(p.proargtypes) WITH ORDINALITY AS arg(type_oid, position)
103 : ORDER BY arg.position
104 : ) END AS argument_types,
105 : p.proargnames AS argument_names,
106 : p.proargmodes AS argument_modes,
107 : COALESCE(pg_catalog.format_type(p.prorettype, NULL), 'any') AS return_type,
108 : p.proretset AS returns_set,
109 : p.prokind AS function_kind
110 : FROM pg_proc p
111 : JOIN pg_namespace n ON n.oid = p.pronamespace
112 : WHERE p.prokind IN ('f', 'a', 'w', 'p')
113 : ORDER BY n.nspname, p.proname
114 : `;
115 :
116 : const args: QueryExecutionArgs = {
117 1 : query: query,
118 1 : abortSignal,
119 1 : readTimeoutMs: CATALOG_QUERY_READ_TIMEOUT_MS,
120 1 : throwOnError: true,
121 1 : metadata: {
122 1 : queryType: QueryType.CATALOG_QUERY_PG_PROC,
123 1 : title: "Query Postgres Functions",
124 1 : description: null,
125 1 : issuer: "Catalog Update",
126 1 : userProvided: false
127 1 : }
128 : };
129 1 : const [queryId, queryExecution] = executor(connectionId, args);
130 1 : connectionDispatch(connectionId, {
131 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
132 1 : value: [updateId, queryId]
133 1 : });
134 :
135 1 : const queryResult = await queryExecution as PgProcTable;
136 1 : return queryResult;
137 1 : }
|
