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 { 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 : /// Pick the statements of a connection script that define or drop relations of user schemas.
23 : /// The text is parsed in a scratch catalog, the shared relation script stays untouched.
24 1 : function describeConnectionRelations(instance: core.DashQL, relationScript: string): ConnectionRelations {
25 1 : const relations: ConnectionRelations = { statements: [], systemRelations: 0 };
26 1 : const catalog = instance.createCatalog();
27 1 : const script = instance.createScript(catalog);
28 1 : let parsed: core.FlatBufferPtr<core.buffers.parser.ParsedScript> | null = null;
29 1 : try {
30 1 : script.replaceText(relationScript);
31 1 : script.parse();
32 1 : parsed = script.getParsed();
33 1 : const reader = parsed.read();
34 1 : const statement = new core.buffers.parser.Statement();
35 1 : for (let i = 0; i < reader.statementsLength(); ++i) {
36 1 : reader.statements(i, statement);
37 1 : const schemaName = statement.target()?.schemaName() ?? '';
38 1 : if (isSystemSchema(schemaName)) {
39 1 : relations.systemRelations += 1;
40 1 : continue;
41 1 : }
42 1 : const span = statement.statementSpan();
43 1 : if (span != null) {
44 1 : relations.statements.push(`${script.toString(span.offset(), span.length())};`);
45 : }
46 1 : }
47 1 : } finally {
48 1 : script.destroy();
49 1 : catalog.destroy();
50 1 : }
51 1 : return relations;
52 1 : }
53 :
54 : /// Render the tables that the session scripts on a connection define as CREATE TABLE statements without types.
55 : /// Queries only see the tables of scripts on their own connection.
56 1 : function describeSessionTables(sessionScripts: SessionScripts, connectionId: string): string[] {
57 1 : const statements: string[] = [];
58 1 : for (const script of Object.values(sessionScripts.scripts)) {
59 1 : const catalog = script.connection.catalog;
60 1 : if (script.connection.connectionId !== connectionId || catalog == null) continue;
61 1 : let entries: core.FlatBufferPtr<core.buffers.catalog.CatalogEntries>;
62 1 : try {
63 1 : entries = catalog.describeEntriesOf(script.scriptSession.getCatalogEntryId());
64 1 : } catch {
65 : // The script is not loaded into the catalog
66 1 : continue;
67 1 : }
68 1 : const reader = entries.read();
69 1 : for (let i = 0; i < reader.entriesLength(); ++i) {
70 1 : const entry = reader.entries(i)!;
71 1 : for (let j = 0; j < entry.schemasLength(); ++j) {
72 1 : const schema = entry.schemas(j)!;
73 1 : for (let k = 0; k < schema.tablesLength(); ++k) {
74 1 : const table = schema.tables(k)!;
75 1 : const columns: string[] = [];
76 1 : for (let l = 0; l < table.columnsLength(); ++l) {
77 1 : columns.push(quoteIdentifier(table.columns(l)!.columnName() ?? ''));
78 : }
79 1 : statements.push(`CREATE TABLE ${quoteIdentifier(table.tableName() ?? '')} (${columns.join(', ')});`);
80 1 : }
81 1 : }
82 1 : }
83 1 : }
84 1 : return statements;
85 1 : }
86 :
87 : /// Build the schema that the model may query on a connection: the user relations of the connection, the catalog
88 : /// changes of executed statements and the tables of the session scripts on it.
89 : /// The schema is cut off at a statement boundary once it exceeds the budget.
90 1 : export function buildSchemaContext(connectionId: string, conn: ConnectionState | null, sessionScripts: SessionScripts, logger: Logger): string {
91 1 : const describe = (text: string): ConnectionRelations => text.trim().length > 0
92 1 : ? describeConnectionRelations(sessionScripts.instance, text)
93 1 : : { statements: [], systemRelations: 0 };
94 1 : const relations = describe(conn?.catalogRelationScript.toString() ?? '');
95 1 : const changes = describe(conn?.catalogChangesSql ?? '');
96 1 : const sessionTables = describeSessionTables(sessionScripts, connectionId);
97 :
98 1 : const statements = [...relations.statements];
99 1 : if (changes.statements.length > 0) {
100 1 : statements.push('-- Changes of executed statements since the last catalog refresh, column types are unknown', ...changes.statements);
101 : }
102 1 : if (sessionTables.length > 0) {
103 1 : statements.push('-- Tables defined by the session, column types are unknown', ...sessionTables);
104 : }
105 1 : let schema = '';
106 1 : let included = 0;
107 1 : for (const statement of statements) {
108 1 : if (schema.length + statement.length + 1 > SCHEMA_CONTEXT_BUDGET) {
109 1 : break;
110 : }
111 1 : schema += schema.length > 0 ? `\n${statement}` : statement;
112 1 : ++included;
113 1 : }
114 1 : const truncated = included < statements.length;
115 1 : if (truncated) {
116 1 : schema += '\n-- The schema is truncated';
117 1 : logger.warn('Truncated the schema for the agent', {
118 1 : budget: String(SCHEMA_CONTEXT_BUDGET),
119 1 : included: String(included),
120 1 : total: String(statements.length),
121 1 : }, LOG_CTX);
122 1 : }
123 1 : logger.debug('Built the schema for the agent', {
124 1 : connectionRelations: String(relations.statements.length),
125 1 : catalogChanges: String(changes.statements.length),
126 1 : skippedSystemRelations: String(relations.systemRelations + changes.systemRelations),
127 1 : sessionTables: String(sessionTables.length),
128 1 : characters: String(schema.length),
129 1 : truncated: String(truncated),
130 1 : }, LOG_CTX);
131 1 : return schema;
132 1 : }
|
