Line data    Source code

       1              : import type * as arrow from 'apache-arrow';
       2              : import type * as connection from '@ankoh/dashql-jsonschema/connection.js';
       3              : 
       4              : import type { InformationSchemaScope } from '../catalog/catalog_query_information_schema.js';
       5              : import type { QueryExecutor } from '../query_executor.js';
       6              : import type { QueryExecutionArgs } from '../query_execution_args.js';
       7              : import type { DynamicConnectionDispatch } from '../connection_registry.js';
       8              : import { CATALOG_UPDATE_REGISTER_QUERY } from '../connection_state.js';
       9              : import { QueryType } from '../../query/query_execution_state.js';
      10              : 
      11              : /// The fields of the Trino params that name catalogs
      12              : type TrinoCatalogParams = Pick<connection.TrinoConnectionParams, 'defaultCatalog' | 'defaultSchema' | 'catalogs' | 'catalogName' | 'schemaNames'>;
      13              : 
      14              : /// The catalog and schema of the names that statements leave out, sent as `X-Trino-Catalog` and `X-Trino-Schema`
      15              : export interface TrinoSessionDefaults {
      16              :     /// The default catalog, null if there is none
      17              :     catalog: string | null;
      18              :     /// The default schema, null if there is none. Trino takes a default schema only with a default catalog.
      19              :     schema: string | null;
      20              : }
      21              : 
      22              : /// A catalog entry `catalog` or `catalog.schema`, split at the first dot. Null if the entry names no catalog.
      23            1 : export function parseTrinoCatalogEntry(entry: string): { catalog: string; schema: string | null } | null {
      24            1 :     const dot = entry.indexOf('.');
      25            1 :     const catalog = (dot < 0 ? entry : entry.substring(0, dot)).trim();
      26            1 :     const schema = dot < 0 ? '' : entry.substring(dot + 1).trim();
      27            1 :     return catalog.length === 0 ? null : { catalog, schema: schema.length === 0 ? null : schema };
      28            1 : }
      29              : 
      30              : /// The params with the fields that replaced `catalogName` and `schemaNames`.
      31              : /// The catalog name was the default catalog, the schema names were the schemas of it that the catalog loader reads.
      32            1 : export function withTrinoCatalogFields<P extends TrinoCatalogParams>(params: P): Omit<P, 'catalogName' | 'schemaNames'> {
      33            1 :     const { catalogName, schemaNames, ...rest } = params;
      34            1 :     const catalog = catalogName ?? '';
      35            1 :     return {
      36            1 :         ...rest,
      37            1 :         defaultCatalog: rest.defaultCatalog ?? catalog,
      38            1 :         catalogs: rest.catalogs ?? (catalog.length === 0 ? [] : (schemaNames ?? []).map(schema => `${catalog}.${schema}`)),
      39            1 :     };
      40            1 : }
      41              : 
      42              : /// The default catalog and schema of the params
      43            1 : export function trinoSessionDefaults(params: TrinoCatalogParams): TrinoSessionDefaults {
      44            1 :     const fields = withTrinoCatalogFields(params);
      45            1 :     const catalog = fields.defaultCatalog?.trim() ?? '';
      46            1 :     const schema = fields.defaultSchema?.trim() ?? '';
      47            1 :     return catalog.length === 0
      48            1 :         ? { catalog: null, schema: null }
      49            1 :         : { catalog, schema: schema.length === 0 ? null : schema };
      50            1 : }
      51              : 
      52              : /// The catalogs that the catalog loader reads, in the order of their first entry.
      53              : /// An entry without a schema reads all schemas of its catalog. Without entries it reads the default catalog.
      54            1 : export function trinoCatalogScopes(params: TrinoCatalogParams): InformationSchemaScope[] {
      55            1 :     const scopes = new Map<string, { schemas: Set<string>; all: boolean }>();
      56            1 :     for (const entry of withTrinoCatalogFields(params).catalogs ?? []) {
      57            1 :         const parsed = parseTrinoCatalogEntry(entry);
      58            1 :         if (parsed == null) {
      59            1 :             continue;
      60              :         }
      61            1 :         let scope = scopes.get(parsed.catalog);
      62            1 :         if (scope == null) {
      63            1 :             scope = { schemas: new Set(), all: false };
      64            1 :             scopes.set(parsed.catalog, scope);
      65            1 :         }
      66            1 :         if (parsed.schema == null) {
      67            1 :             scope.all = true;
      68              :         } else {
      69            1 :             scope.schemas.add(parsed.schema);
      70              :         }
      71            1 :     }
      72            1 :     if (scopes.size === 0) {
      73            1 :         const { catalog } = trinoSessionDefaults(params);
      74            1 :         return catalog == null ? [] : [{ catalog, schemas: [] }];
      75            1 :     }
      76            1 :     return [...scopes].map(([catalog, scope]) => ({ catalog, schemas: scope.all ? [] : [...scope.schemas] }));
      77            1 : }
      78              : 
      79              : /// The query of the catalogs of the server
      80            1 : export const TRINO_CATALOGS_QUERY = 'SELECT catalog_name FROM system.metadata.catalogs ORDER BY catalog_name';
      81              : 
      82              : /// The catalogs of the server, all schemas of each.
      83              : /// The catalog loader reads them when the params name no catalogs.
      84            1 : export async function queryTrinoCatalogScopes(connectionId: string, connectionDispatch: DynamicConnectionDispatch, updateId: number, executor: QueryExecutor, abortSignal?: AbortSignal): Promise<InformationSchemaScope[]> {
      85              :     const args: QueryExecutionArgs = {
      86            1 :         query: TRINO_CATALOGS_QUERY,
      87            1 :         abortSignal,
      88            1 :         throwOnError: true,
      89            1 :         metadata: {
      90            1 :             queryType: QueryType.CATALOG_QUERY_INFORMATION_SCHEMA,
      91            1 :             title: "Catalogs",
      92            1 :             description: null,
      93            1 :             issuer: "Catalog Update",
      94            1 :             userProvided: false
      95            1 :         }
      96              :     };
      97            1 :     const [queryId, queryExecution] = executor(connectionId, args);
      98            1 :     connectionDispatch(connectionId, {
      99            1 :         type: CATALOG_UPDATE_REGISTER_QUERY,
     100            1 :         value: [updateId, queryId]
     101            1 :     });
     102            1 :     const result = await queryExecution as arrow.Table<{ catalog_name: arrow.Utf8 }> | null;
     103            1 :     const scopes: InformationSchemaScope[] = [];
     104            1 :     for (const batch of result?.batches ?? []) {
     105            1 :         const names = batch.getChild('catalog_name');
     106            1 :         for (let i = 0; i < batch.numRows; ++i) {
     107            1 :             const name = names?.at(i);
     108            1 :             if (name) {
     109            1 :                 scopes.push({ catalog: name, schemas: [] });
     110              :             }
     111            1 :         }
     112            1 :     }
     113            1 :     return scopes;
     114            1 : }