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 { catalogTablesOfSchema, quoteIdentifier, relationFilter, type CatalogTable, type ColumnMetadata, type RelationName } from './catalog_sql_generator.js';
9 :
10 1 : export const CATALOG_QUERY_READ_TIMEOUT_MS = 60_000;
11 :
12 : export type PgAttributeColumnsTable = arrow.Table<{
13 : table_schema: arrow.Utf8;
14 : table_name: arrow.Utf8;
15 : column_name: arrow.Utf8;
16 : ordinal_position: arrow.Int32;
17 : data_type: arrow.Utf8;
18 : }>;
19 :
20 : /**
21 : * The tables of pg_attribute query results.
22 : * Builds a hierarchy of schema -> table -> columns.
23 : */
24 1 : export function catalogTablesFromPgAttribute(result: PgAttributeColumnsTable, databaseName: string): CatalogTable[] {
25 : // Build hierarchy: schema -> table -> columns
26 1 : const schemas = new Map<string, Map<string, ColumnMetadata[]>>();
27 :
28 1 : for (const batch of result.batches) {
29 1 : const colTableSchema = batch.getChild("table_schema")!;
30 1 : const colTableName = batch.getChild("table_name")!;
31 1 : const colColumnName = batch.getChild("column_name")!;
32 1 : const colOrdinalPos = batch.getChild("ordinal_position")!;
33 1 : const colDataType = batch.getChild("data_type")!;
34 :
35 : // Iterate over all rows in the batch
36 1 : for (let i = 0; i < batch.numRows; ++i) {
37 1 : const tableSchema = colTableSchema.at(i);
38 1 : const tableName = colTableName.at(i);
39 1 : const columnName = colColumnName.at(i);
40 1 : const ordinalPosition = colOrdinalPos.at(i);
41 1 : const dataType = colDataType.at(i);
42 :
43 1 : if (!tableSchema || !columnName || !tableName) {
44 1 : continue;
45 : }
46 :
47 : // Get or create schema entry
48 1 : let tableMap = schemas.get(tableSchema);
49 1 : if (!tableMap) {
50 1 : tableMap = new Map<string, ColumnMetadata[]>();
51 1 : schemas.set(tableSchema, tableMap);
52 1 : }
53 :
54 : // Get or create table entry
55 1 : let columns = tableMap.get(tableName);
56 1 : if (!columns) {
57 1 : columns = [];
58 1 : tableMap.set(tableName, columns);
59 1 : }
60 :
61 : // Add column metadata
62 1 : columns.push({
63 1 : name: columnName,
64 1 : ordinalPosition: ordinalPosition ?? columns.length,
65 1 : dataType: dataType ?? null,
66 1 : });
67 1 : }
68 1 : }
69 :
70 1 : return [...schemas].flatMap(([schemaName, tables]) => catalogTablesOfSchema(databaseName, schemaName, tables));
71 1 : }
72 :
73 : /// What a query of pg_attribute reads
74 : export interface PgAttributeScope {
75 : /// The database to read the catalog of, the default database if absent
76 : queryDatabaseName?: string | null;
77 : /// The relations, all relations if absent
78 : relations?: RelationName[] | null;
79 : }
80 :
81 : /// The query of the columns of relations
82 1 : export function pgAttributeQuery(scope: PgAttributeScope = {}): string {
83 1 : const pgCatalog = scope.queryDatabaseName == null
84 1 : ? 'pg_catalog'
85 1 : : `${quoteIdentifier(scope.queryDatabaseName)}.${quoteIdentifier('pg_catalog')}`;
86 1 : const relationsFilter = scope.relations != null ? `AND ${relationFilter('n.nspname', 'c.relname', scope.relations)}` : '';
87 : // A name without a schema doesn't mean a table of the system
88 1 : const systemFilter = scope.relations != null ? `AND n.nspname NOT IN ('pg_catalog', 'information_schema')` : '';
89 1 : return `
90 : SELECT
91 : n.nspname AS table_schema,
92 : c.relname AS table_name,
93 : a.attname AS column_name,
94 : a.attnum AS ordinal_position,
95 : t.typname AS data_type
96 1 : FROM ${pgCatalog}.pg_class c
97 1 : JOIN ${pgCatalog}.pg_namespace n ON n.oid = c.relnamespace
98 1 : JOIN ${pgCatalog}.pg_attribute a ON a.attrelid = c.oid
99 1 : JOIN ${pgCatalog}.pg_type t ON t.oid = a.atttypid
100 : WHERE c.relkind IN ('r', 'v', 'm', 'f', 'p')
101 : AND a.attnum > 0
102 : AND NOT a.attisdropped
103 1 : ${relationsFilter}
104 1 : ${systemFilter}
105 : ORDER BY n.nspname, c.relname, a.attnum
106 : `;
107 1 : }
108 :
109 1 : export async function queryPgAttribute(
110 1 : connectionId: string,
111 1 : connectionDispatch: DynamicConnectionDispatch,
112 1 : updateId: number,
113 1 : executor: QueryExecutor,
114 1 : queryDatabaseName: string | null = null,
115 1 : abortSignal?: AbortSignal,
116 1 : relations: RelationName[] | null = null,
117 1 : ): Promise<PgAttributeColumnsTable | null> {
118 1 : const query = pgAttributeQuery({ queryDatabaseName, relations });
119 :
120 : const args: QueryExecutionArgs = {
121 1 : query: query,
122 1 : abortSignal,
123 1 : readTimeoutMs: CATALOG_QUERY_READ_TIMEOUT_MS,
124 1 : throwOnError: true,
125 1 : metadata: {
126 1 : queryType: QueryType.CATALOG_QUERY_PG_ATTRIBUTE,
127 1 : title: "Query Postgres Schema Columns",
128 1 : description: null,
129 1 : issuer: "Catalog Update",
130 1 : userProvided: false
131 1 : }
132 : };
133 1 : const [queryId, queryExecution] = executor(connectionId, args);
134 1 : connectionDispatch(connectionId, {
135 1 : type: CATALOG_UPDATE_REGISTER_QUERY,
136 1 : value: [updateId, queryId]
137 1 : });
138 :
139 1 : const queryResult = await queryExecution as PgAttributeColumnsTable;
140 1 : return queryResult;
141 1 : }
|
