Line data    Source code

       1              : import * as core from '@ankoh/dashql-core';
       2              : 
       3              : import type { Logger } from '../platform/logger/logger.js';
       4              : import type { ConnectionState } from '../connections/connection_state.js';
       5              : import type { SessionScripts } from '../session/scripts/session_scripts.js';
       6              : import { catalogTablesOfDescriptor, generateCatalogTablesSQL, quoteIdentifier } from '../connections/catalog/catalog_sql_generator.js';
       7              : import { isSystemSchema } from '../connections/catalog/system_schemas.js';
       8              : 
       9            1 : const LOG_CTX = 'agent_catalog';
      10              : 
      11              : /// The schema context is cut off after this many characters to keep the output of `dashql agent catalog` readable
      12            1 : const SCHEMA_CONTEXT_BUDGET = 100_000;
      13              : 
      14              : /// The relations of a connection, split by owner
      15              : interface ConnectionRelations {
      16              :     /// The statements that define relations of user schemas
      17              :     statements: string[];
      18              :     /// The number of skipped relations of system schemas
      19              :     systemRelations: number;
      20              : }
      21              : 
      22              : /// The statements of the relations of user schemas of a connection
      23            1 : function describeConnectionRelations(descriptor: core.buffers.catalog.CatalogDescriptorT | null): ConnectionRelations {
      24            1 :     const relations: ConnectionRelations = { statements: [], systemRelations: 0 };
      25            1 :     for (const table of catalogTablesOfDescriptor(descriptor)) {
      26            1 :         if (isSystemSchema(table.schemaName)) {
      27            1 :             relations.systemRelations += 1;
      28            1 :             continue;
      29            1 :         }
      30            1 :         relations.statements.push(generateCatalogTablesSQL([table]));
      31            1 :     }
      32            1 :     return relations;
      33            1 : }
      34              : 
      35              : /// Render the tables that the session scripts on a connection define as CREATE TABLE statements without types.
      36              : /// Queries only see the tables of scripts on their own connection.
      37            1 : function describeSessionTables(sessionScripts: SessionScripts, connectionId: string): string[] {
      38            1 :     const statements: string[] = [];
      39            1 :     for (const script of Object.values(sessionScripts.scripts)) {
      40            1 :         const catalog = script.connection.catalog;
      41            1 :         if (script.connection.connectionId !== connectionId || catalog == null) continue;
      42            1 :         let entries: core.FlatBufferPtr<core.buffers.catalog.CatalogEntries>;
      43            1 :         try {
      44            1 :             entries = catalog.describeEntriesOf(script.scriptSession.getCatalogEntryId());
      45            1 :         } catch {
      46              :             // The script is not loaded into the catalog
      47            1 :             continue;
      48            1 :         }
      49            1 :         const reader = entries.read();
      50            1 :         for (let i = 0; i < reader.entriesLength(); ++i) {
      51            1 :             const entry = reader.entries(i)!;
      52            1 :             for (let j = 0; j < entry.schemasLength(); ++j) {
      53            1 :                 const schema = entry.schemas(j)!;
      54            1 :                 for (let k = 0; k < schema.tablesLength(); ++k) {
      55            1 :                     const table = schema.tables(k)!;
      56            1 :                     const columns: string[] = [];
      57            1 :                     for (let l = 0; l < table.columnsLength(); ++l) {
      58            1 :                         columns.push(quoteIdentifier(table.columns(l)!.columnName() ?? ''));
      59              :                     }
      60            1 :                     statements.push(`CREATE TABLE ${quoteIdentifier(table.tableName() ?? '')} (${columns.join(', ')});`);
      61            1 :                 }
      62            1 :             }
      63            1 :         }
      64            1 :     }
      65            1 :     return statements;
      66            1 : }
      67              : 
      68              : /// Build the schema that the model may query on a connection: the user relations of the connection and the tables of
      69              : /// the session scripts on it.
      70              : /// The schema is cut off at a statement boundary once it exceeds the budget.
      71            1 : export function buildSchemaContext(connectionId: string, conn: ConnectionState | null, sessionScripts: SessionScripts, logger: Logger): string {
      72            1 :     const relations = describeConnectionRelations(conn?.catalog.readDescriptor(conn.catalogRelationEntryId) ?? null);
      73            1 :     const sessionTables = describeSessionTables(sessionScripts, connectionId);
      74              : 
      75            1 :     const statements = [...relations.statements];
      76            1 :     if (sessionTables.length > 0) {
      77            1 :         statements.push('-- Tables defined by the session, column types are unknown', ...sessionTables);
      78              :     }
      79            1 :     let schema = '';
      80            1 :     let included = 0;
      81            1 :     for (const statement of statements) {
      82            1 :         if (schema.length + statement.length + 1 > SCHEMA_CONTEXT_BUDGET) {
      83            1 :             break;
      84              :         }
      85            1 :         schema += schema.length > 0 ? `\n${statement}` : statement;
      86            1 :         ++included;
      87            1 :     }
      88            1 :     const truncated = included < statements.length;
      89            1 :     if (truncated) {
      90            1 :         schema += '\n-- The schema is truncated';
      91            1 :         logger.warn('Truncated the schema for the agent', {
      92            1 :             budget: String(SCHEMA_CONTEXT_BUDGET),
      93            1 :             included: String(included),
      94            1 :             total: String(statements.length),
      95            1 :         }, LOG_CTX);
      96            1 :     }
      97            1 :     logger.debug('Built the schema for the agent', {
      98            1 :         connectionRelations: String(relations.statements.length),
      99            1 :         skippedSystemRelations: String(relations.systemRelations),
     100            1 :         sessionTables: String(sessionTables.length),
     101            1 :         characters: String(schema.length),
     102            1 :         truncated: String(truncated),
     103            1 :     }, LOG_CTX);
     104            1 :     return schema;
     105            1 : }