Line data Source code
1 : import type * as arrow from 'apache-arrow';
2 : import * as core from '@ankoh/dashql-core';
3 :
4 : import type { LoggerLike, Logger } from '../platform/logger/logger.js';
5 : import { stringifyError } from '../platform/logger/logger.js';
6 : import type { ClientFunctionSpec } from '../connections/clfn/clfn_spec.js';
7 : import { connectorOfDetails } from '../connections/connector_descriptor.js';
8 : import type { AgentCommandParams, AgentCommandResult } from './agent_protocol.js';
9 : import { AgentCommandError } from './agent_client.js';
10 : import type { ConnectionRegistry } from '../connections/connection_registry.js';
11 : import { ConnectionHealth, connectionCachesQueryResults, type ConnectionState } from '../connections/connection_state.js';
12 : import { isHyperWasmConnection } from '../connections/connection_label.js';
13 : import { findEmbeddedConnection } from '../connections/embedded_connection.js';
14 : import type { QueryExecutor } from '../connections/query_executor.js';
15 : import {
16 : CREATE_SCRIPT,
17 : RENAME_SCRIPT,
18 : SET_SCRIPT_CONNECTION,
19 : SET_SCRIPT_TEXT,
20 : getSelectedScriptRef,
21 : newScriptConnection,
22 : persistedConnectionId,
23 : scriptConnectionOf,
24 : sessionConnectionIds,
25 : type ScriptConnection,
26 : type ScriptData,
27 : type SessionScripts,
28 : type SessionScriptsAction,
29 : } from '../session/scripts/session_scripts.js';
30 : import { ensureSessionScriptAnalyzed, type ModifySessionScripts } from '../session/scripts/session_scripts_registry.js';
31 : import { scriptDisplayName } from '../session/scripts/script_types.js';
32 : import { runSessionScriptForResult } from '../session/scripts/rerun_query.js';
33 : import { buildSchemaContext } from './agent_catalog.js';
34 : import { checkVegaLiteSpec } from '../ui/vegalite/vegalite_check.js';
35 : import { arrowTableToRows } from '../ui/vegalite/vegalite_rows.js';
36 : import { CHART_IMAGE_HEIGHT, CHART_IMAGE_SCALE, CHART_IMAGE_WIDTH, renderVegaLitePng, type ChartImage } from '../ui/vegalite/vegalite_image.js';
37 : import { DEFAULT_RUN_ROWS, RESULT_PAGE_ROWS, formatResultSlice, sliceResult } from './agent_result_format.js';
38 : import { formatBytes } from '../utils/format.js';
39 :
40 : /// What the commands of the agent work with
41 : export interface AgentCommandEnvironment {
42 : /// The scripts of a session, null if the session is gone
43 : sessionScripts: (sessionId: string) => SessionScripts | null;
44 : /// Modify the scripts of a session
45 : modifySessionScripts: (sessionId: string, action: SessionScriptsAction) => Promise<SessionScripts | null>;
46 : /// The connections of the app
47 : connections: () => ConnectionRegistry;
48 : /// The connection to the embedded Hyper, allocated when the app holds none yet
49 : embeddedConnection?: (instance: core.DashQL) => ConnectionState;
50 : executeQuery: QueryExecutor;
51 : logger: Logger;
52 : /// The agent edited a script
53 : onScriptEdited?: (sessionId: string, scriptKey: number) => void;
54 : /// The agent ran a script
55 : onScriptRun?: (sessionId: string, scriptKey: number) => void;
56 : }
57 :
58 1 : function fail(message: string): never {
59 1 : throw new AgentCommandError('failed', message);
60 1 : }
61 :
62 1 : function usage(message: string): never {
63 1 : throw new AgentCommandError('usage', message);
64 1 : }
65 :
66 1 : function stringArg(args: Record<string, unknown>, name: string): string {
67 1 : const value = args[name];
68 1 : if (typeof value !== 'string' || value.length === 0) usage(`missing ${name}`);
69 1 : return value;
70 1 : }
71 :
72 1 : function numberArg(args: Record<string, unknown>, name: string, fallback: number): number {
73 1 : const value = args[name];
74 1 : if (value == null) return fallback;
75 1 : if (typeof value !== 'number' || !Number.isInteger(value) || value <= 0) usage(`${name} must be a positive integer`);
76 1 : return value;
77 1 : }
78 :
79 1 : function booleanArg(args: Record<string, unknown>, name: string): boolean {
80 1 : const value = args[name];
81 1 : if (value == null) return false;
82 1 : if (typeof value !== 'boolean') usage(`${name} must be a boolean`);
83 1 : return value;
84 1 : }
85 :
86 : /// Find a script by its file name or display name, case-insensitive
87 1 : export function findScript(state: SessionScripts, name: string): ScriptData | null {
88 1 : const wanted = scriptDisplayName(name.trim()).toLowerCase();
89 1 : for (const fileName of state.scriptOrder) {
90 1 : if (scriptDisplayName(fileName).toLowerCase() === wanted) {
91 1 : return state.scripts[state.scriptRefs[fileName].scriptId] ?? null;
92 : }
93 : }
94 1 : return null;
95 1 : }
96 :
97 1 : function requireScript(state: SessionScripts, name: string): ScriptData {
98 1 : const script = findScript(state, name);
99 1 : if (script == null) {
100 1 : const names = state.scriptOrder.map(scriptDisplayName).join(', ');
101 1 : fail(`No script named ${name}. Scripts: ${names.length > 0 ? names : 'none'}`);
102 1 : }
103 1 : return script;
104 1 : }
105 :
106 : /// A connection of a session as the agent sees it
107 : interface AgentConnection {
108 : id: string;
109 : connector: string | null;
110 : /// Is it the embedded Hyper, which reads the attachments and the results of scripts of other connections?
111 : embedded: boolean;
112 : readOnly: boolean | null;
113 : online: boolean;
114 : /// Does DashQL cache the query results of the connection, so the embedded Hyper rereads them without a query?
115 : cachesResults: boolean | null;
116 : }
117 :
118 1 : function describeConnection(env: AgentCommandEnvironment, connectionId: string): AgentConnection {
119 1 : const conn = env.connections().connections.get(connectionId);
120 1 : if (conn == null) return { id: connectionId, connector: null, embedded: false, readOnly: null, online: false, cachesResults: null };
121 1 : const embedded = isHyperWasmConnection(conn);
122 1 : return {
123 1 : id: connectionId,
124 1 : connector: conn.connectorInfo.names.displayShort,
125 1 : embedded,
126 1 : readOnly: conn.connectorInfo.features.readOnly,
127 1 : online: conn.connectionHealth === ConnectionHealth.ONLINE,
128 : // The embedded Hyper reads its own tables directly, its cache saves no remote query
129 1 : cachesResults: !embedded && connectionCachesQueryResults(conn),
130 1 : };
131 1 : }
132 :
133 : /// The connections that the agent may use: the connections of the session and the embedded Hyper, which joins the
134 : /// attachments and the results of the other connections. The app may hold no embedded Hyper yet, e.g. when every
135 : /// session runs on a remote connection, then it is allocated so the agent can move scripts onto it.
136 1 : function agentConnectionIds(env: AgentCommandEnvironment, state: SessionScripts): string[] {
137 1 : const ids = sessionConnectionIds(state);
138 1 : const embedded = findEmbeddedConnection(env.connections().connections.values()) ?? env.embeddedConnection?.(state.instance) ?? null;
139 1 : if (embedded != null && !ids.includes(embedded.connectionId)) ids.push(embedded.connectionId);
140 1 : return ids;
141 1 : }
142 :
143 : /// A connection that the agent may use
144 1 : function requireSessionConnection(env: AgentCommandEnvironment, state: SessionScripts, connectionId: string): void {
145 1 : const ids = agentConnectionIds(env, state);
146 1 : if (!ids.includes(connectionId)) fail(`The session does not use the connection ${connectionId}. Connections: ${ids.join(', ')}`);
147 1 : }
148 :
149 : /// The connection that the agent picked for a script, null if it picked none
150 1 : function pickedScriptConnection(env: AgentCommandEnvironment, state: SessionScripts, args: Record<string, unknown>): ScriptConnection | null {
151 1 : const connectionId = args.connection;
152 1 : if (connectionId == null) return null;
153 1 : if (typeof connectionId !== 'string' || connectionId.length === 0) usage('invalid connection');
154 1 : requireSessionConnection(env, state, connectionId);
155 1 : const conn = env.connections().connections.get(connectionId);
156 1 : if (conn == null) fail(`The connection ${connectionId} is not set up in DashQL, the user has to connect it`);
157 1 : return scriptConnectionOf(conn);
158 1 : }
159 :
160 1 : function connectionLabel(conn: AgentConnection): string {
161 1 : if (conn.connector == null) return conn.id;
162 1 : const connector = conn.embedded ? 'embedded Hyper' : conn.connector;
163 1 : const caching = conn.cachesResults ? ', caches results' : '';
164 1 : return `${conn.id} (${connector}, ${conn.readOnly ? 'read-only' : 'read-write'}, ${conn.online ? 'online' : 'offline'}${caching})`;
165 1 : }
166 :
167 : /// The SHA-256 of a text as hex
168 1 : async function sha256(text: string): Promise<string> {
169 1 : const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(text));
170 1 : return [...new Uint8Array(digest)].map(b => b.toString(16).padStart(2, '0')).join('');
171 1 : }
172 :
173 : /// The line and column of an offset, one-based
174 1 : function lineAndColumn(text: string, offset: number): [number, number] {
175 1 : let line = 1;
176 1 : let lineStart = 0;
177 1 : for (let i = 0; i < Math.min(offset, text.length); ++i) {
178 1 : if (text.charCodeAt(i) === 10) {
179 0 : line += 1;
180 0 : lineStart = i + 1;
181 0 : }
182 : }
183 1 : return [line, offset - lineStart + 1];
184 1 : }
185 :
186 : /// A diagnostic of a script
187 : interface AgentDiagnostic {
188 : severity: 'error' | 'warning';
189 : line: number | null;
190 : column: number | null;
191 : message: string;
192 : }
193 :
194 1 : function collectDiagnostics(script: ScriptData): AgentDiagnostic[] {
195 1 : const text = script.scriptSession.getText();
196 1 : return (script.editorUpdate?.diagnostics ?? []).map(diagnostic => {
197 1 : const span = diagnostic.textSpan;
198 1 : const [line, column] = span == null ? [null, null] : lineAndColumn(text, Number(span.offset));
199 1 : return {
200 1 : severity: diagnostic.severity === core.buffers.editor.EditorDiagnosticSeverity.WARNING ? 'warning' : 'error',
201 1 : line,
202 1 : column,
203 1 : message: typeof diagnostic.message === 'string' ? diagnostic.message : '',
204 1 : };
205 1 : });
206 1 : }
207 :
208 1 : function formatDiagnostics(name: string, diagnostics: AgentDiagnostic[]): string {
209 1 : if (diagnostics.length === 0) return `${name}: no diagnostics`;
210 1 : return diagnostics.map(d => `${name}:${d.line ?? '?'}:${d.column ?? '?'}: ${d.severity}: ${d.message}`).join('\n');
211 1 : }
212 :
213 : /// Analyze a script if its analysis is outdated
214 1 : async function analyzed(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, script: ScriptData): Promise<ScriptData> {
215 1 : const modify: ModifySessionScripts = action => env.modifySessionScripts(sessionId, action);
216 1 : return (await ensureSessionScriptAnalyzed(state, script.scriptKey, modify)) ?? script;
217 1 : }
218 :
219 1 : async function listScripts(env: AgentCommandEnvironment, state: SessionScripts): Promise<AgentCommandResult> {
220 1 : const focused = getSelectedScriptRef(state)?.fileName ?? null;
221 1 : const scripts = state.scriptOrder.map(fileName => {
222 1 : const script = state.scripts[state.scriptRefs[fileName].scriptId];
223 1 : const text = script?.scriptSession.getText() ?? '';
224 1 : return {
225 1 : name: scriptDisplayName(fileName),
226 1 : fileName,
227 1 : connectionId: script == null ? null : persistedConnectionId(script),
228 1 : lines: text.length === 0 ? 0 : text.split('\n').length,
229 1 : focused: fileName === focused,
230 1 : };
231 1 : });
232 1 : const connections = agentConnectionIds(env, state).map(id => describeConnection(env, id));
233 1 : const attachments = (state.attachments ?? []).map(attachment => ({
234 1 : fileName: attachment.fileName,
235 1 : name: attachment.name,
236 1 : sizeBytes: attachment.sizeBytes,
237 1 : columns: attachment.columns,
238 1 : }));
239 1 : const lines = [
240 1 : `Session ${JSON.stringify(state.name ?? '')} (${state.sessionId})`,
241 1 : `Connections: ${connections.map(connectionLabel).join(', ')}`,
242 1 : ];
243 1 : lines.push(scripts.length === 0 ? 'No scripts' : 'Scripts:');
244 : // The scripts name their connections when the session uses several
245 1 : const several = connections.length > 1;
246 1 : for (const script of scripts) {
247 1 : const connection = several && script.connectionId != null ? ` on ${script.connectionId}` : '';
248 1 : lines.push(` ${script.name} ${script.lines} line${script.lines === 1 ? '' : 's'}${connection}${script.focused ? ' (focused)' : ''}`);
249 1 : }
250 : // The attachments are table parameters of the scripts of the embedded Hyper
251 1 : if (attachments.length > 0) lines.push('Attachments:');
252 1 : for (const attachment of attachments) {
253 1 : const what = attachment.columns == null ? `:${attachment.name}` : `:${attachment.name} (${attachment.columns.join(', ')})`;
254 1 : lines.push(` ${attachment.fileName} ${formatBytes(attachment.sizeBytes)} ${what}`);
255 1 : }
256 1 : return { text: lines.join('\n'), data: { session: state.sessionId, name: state.name, connections, scripts, attachments } };
257 1 : }
258 :
259 1 : async function editScript(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
260 1 : const name = stringArg(args, 'script');
261 1 : const text = args.text;
262 1 : if (typeof text !== 'string') usage('missing text');
263 1 : const base = scriptDisplayName(name.trim());
264 1 : if (base.length === 0 || /[/\\\0]/.test(base)) usage(`invalid script name: ${name}`);
265 :
266 1 : const connection = pickedScriptConnection(env, state, args);
267 1 : let script = findScript(state, name);
268 1 : let next: SessionScripts | null;
269 1 : let created = false;
270 1 : if (script != null) {
271 1 : if (connection != null && connection.connectionId !== script.connection.connectionId) {
272 : // The script moves first and keeps its key
273 1 : const moved = await env.modifySessionScripts(sessionId, { type: SET_SCRIPT_CONNECTION, value: { fileName: script.fileName, connection } });
274 1 : if (moved?.scripts[script.scriptKey]?.connection.connectionId !== connection.connectionId) {
275 1 : fail(`Cannot move ${scriptDisplayName(script.fileName)} to the connection ${connection.connectionId}`);
276 : }
277 : }
278 1 : next = await env.modifySessionScripts(sessionId, { type: SET_SCRIPT_TEXT, value: { scriptKey: script.scriptKey, text } });
279 1 : } else {
280 1 : next = await env.modifySessionScripts(sessionId, { type: CREATE_SCRIPT, value: { index: null, text, connection: connection ?? undefined } });
281 1 : const createdRef = next == null ? undefined : getSelectedScriptRef(next);
282 1 : if (next == null || createdRef == null) fail('Cannot create the script');
283 1 : next = await env.modifySessionScripts(sessionId, { type: RENAME_SCRIPT, value: { fileName: createdRef.fileName, newFileName: base } });
284 1 : created = true;
285 1 : script = next == null ? null : next.scripts[createdRef.scriptId] ?? null;
286 1 : }
287 1 : if (next == null || script == null) fail('The session is gone');
288 1 : const edited = next.scripts[script.scriptKey];
289 1 : if (edited == null) fail('The script is gone');
290 1 : return reportEdit(env, sessionId, next, edited, created ? 'Created' : 'Updated', { created });
291 1 : }
292 :
293 : /// Analyze an edited script and report its hash and diagnostics
294 1 : async function reportEdit(
295 1 : env: AgentCommandEnvironment,
296 1 : sessionId: string,
297 1 : state: SessionScripts,
298 1 : edited: ScriptData,
299 1 : verb: string,
300 1 : data: Record<string, unknown>,
301 1 : ): Promise<AgentCommandResult> {
302 1 : env.onScriptEdited?.(sessionId, edited.scriptKey);
303 1 : const checked = await analyzed(env, sessionId, state, edited);
304 1 : const diagnostics = collectDiagnostics(checked);
305 1 : const hash = await sha256(checked.scriptSession.getText());
306 1 : const errors = diagnostics.filter(d => d.severity === 'error').length;
307 1 : const fileName = scriptDisplayName(checked.fileName);
308 1 : const connectionId = checked.connection.connectionId;
309 1 : const on = sessionConnectionIds(state).length > 1 ? ` on ${connectionId}` : '';
310 1 : const lines = [`${verb} ${fileName}${on}: sha256 ${hash}, ${errors} error${errors === 1 ? '' : 's'}, ${diagnostics.length - errors} warning${diagnostics.length - errors === 1 ? '' : 's'}`];
311 1 : if (diagnostics.length > 0) lines.push(formatDiagnostics(fileName, diagnostics));
312 1 : return { text: lines.join('\n'), data: { script: fileName, ...data, connection: connectionId, sha256: hash, diagnostics } };
313 1 : }
314 :
315 : /// Rewrite the `@vis` tag of the terminal query in a scratch script and analyze the result. Returns the
316 : /// rewritten text and the analyzer warnings about the new tag, e.g. fields that are no output columns.
317 1 : function rewriteVisTag(state: SessionScripts, script: ScriptData, spec: string): [string, string[]] {
318 1 : const scratch = state.instance.createScript(script.connection.catalog);
319 1 : try {
320 1 : scratch.replaceText(script.scriptSession.getText());
321 1 : const tag = scratch.rewriteVegaLite(spec);
322 1 : scratch.analyze();
323 1 : const analyzedPtr = scratch.getAnalyzed();
324 1 : const reader = analyzedPtr.read();
325 1 : const problems: string[] = [];
326 1 : for (let i = 0; i < reader.errorsLength(); ++i) {
327 1 : const error = reader.errors(i)!;
328 1 : const span = error.textSpan();
329 1 : if (error.errorType() !== core.buffers.analyzer.AnalyzerErrorType.INVALID_VIS_TAG || span == null) continue;
330 1 : if (span.offset() < tag.offset + tag.length && span.offset() + span.length() > tag.offset) problems.push(error.message() ?? '');
331 1 : }
332 1 : return [scratch.toString(), problems];
333 1 : } finally {
334 1 : scratch.destroy();
335 1 : }
336 1 : }
337 :
338 : /// Visualize the terminal query of a script with a Vega-Lite spec. The spec is checked first: it must be a
339 : /// JSON object that Vega-Lite compiles without warnings, and its fields must be output columns of the query.
340 1 : async function editVis(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
341 1 : const script = requireScript(state, stringArg(args, 'script'));
342 1 : const specText = stringArg(args, 'spec');
343 1 : const name = scriptDisplayName(script.fileName);
344 1 : let spec: unknown;
345 1 : try {
346 1 : spec = JSON.parse(specText);
347 1 : } catch (e: unknown) {
348 1 : fail(`The Vega-Lite spec is no valid JSON: ${stringifyError(e)}`);
349 1 : }
350 1 : if (spec == null || typeof spec !== 'object' || Array.isArray(spec)) fail('The Vega-Lite spec is no JSON object');
351 1 : const invalid = (problems: string[]) => fail([`The Vega-Lite spec is invalid, ${name} is unchanged:`, ...problems.map(p => ` ${p}`)].join('\n'));
352 :
353 1 : const compileProblems = await checkVegaLiteSpec(spec as Record<string, unknown>);
354 1 : if (compileProblems.length > 0) invalid(compileProblems);
355 1 : let after: string;
356 1 : let tagProblems: string[];
357 1 : try {
358 1 : [after, tagProblems] = rewriteVisTag(state, script, specText);
359 1 : } catch (e: unknown) {
360 1 : fail(`Cannot visualize ${name}: ${stringifyError(e)}`);
361 1 : }
362 1 : if (tagProblems.length > 0) invalid(tagProblems);
363 :
364 1 : const next = await env.modifySessionScripts(sessionId, { type: SET_SCRIPT_TEXT, value: { scriptKey: script.scriptKey, text: after } });
365 1 : const edited = next?.scripts[script.scriptKey];
366 1 : if (next == null || edited == null) fail('The script is gone');
367 1 : return reportEdit(env, sessionId, next, edited, 'Visualized', {});
368 1 : }
369 :
370 : /// The client functions of a connection as lines of a schema: a family per line, with what it reads and what it writes
371 1 : async function describeClientFunctions(env: AgentCommandEnvironment, conn: ConnectionState | null): Promise<string[]> {
372 1 : const binding = conn == null ? undefined : connectorOfDetails(conn.details).clientFunctions;
373 1 : if (binding == null) return [];
374 1 : let spec: ClientFunctionSpec;
375 1 : try {
376 1 : spec = await binding.loadSpec();
377 1 : } catch (e: unknown) {
378 0 : env.logger.warn('Failed to load the client functions', { error: stringifyError(e) }, 'agent_commands');
379 0 : return [];
380 0 : }
381 1 : const families = new Map<string, { reads: string[]; writes: string[] }>();
382 1 : for (const fn of spec.functions) {
383 1 : const family = fn.family ?? spec.namespace;
384 1 : const names = families.get(family) ?? { reads: [], writes: [] };
385 1 : (fn.effect === 'read' ? names.reads : names.writes).push(fn.name);
386 1 : families.set(family, names);
387 1 : }
388 1 : const lines = [...families].map(([family, { reads, writes }]) => `-- ${family}: ${[reads.join(', '), writes.length > 0 ? `writes: ${writes.join(', ')}` : ''].filter(s => s.length > 0).join(' | ')}`);
389 1 : return [`-- The client functions of the connection, call them with CALL ${spec.namespace}.<function> (<option> => <value>, ...)`, ...lines];
390 1 : }
391 :
392 1 : async function describeCatalog(env: AgentCommandEnvironment, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
393 1 : const connectionId = typeof args.connection === 'string' ? args.connection : newScriptConnection(state, null)?.connectionId;
394 1 : if (connectionId == null) fail('The session has no connection');
395 1 : requireSessionConnection(env, state, connectionId);
396 1 : const conn: ConnectionState | null = env.connections().connections.get(connectionId) ?? null;
397 1 : let schema = buildSchemaContext(connectionId, conn, state, env.logger);
398 : // The embedded Hyper reads the Parquet attachments as table parameters
399 1 : const attachments = conn != null && isHyperWasmConnection(conn) ? state.attachments ?? [] : [];
400 1 : if (attachments.length > 0) {
401 1 : const lines = attachments.map(a => `-- :${a.name}${a.columns == null ? '' : ` (${a.columns.join(', ')})`}`);
402 1 : schema = [schema, '-- The attachments of the session, read them as table parameters', ...lines].filter(s => s.length > 0).join('\n');
403 1 : }
404 1 : schema = [schema, ...await describeClientFunctions(env, conn)].filter(s => s.length > 0).join('\n');
405 1 : const text = schema.length > 0 ? schema : '-- The catalog is empty';
406 1 : return { text, data: { connection: connectionId, schema } };
407 1 : }
408 :
409 1 : async function lintScript(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
410 1 : const script = await analyzed(env, sessionId, state, requireScript(state, stringArg(args, 'script')));
411 1 : const diagnostics = collectDiagnostics(script);
412 1 : const name = scriptDisplayName(script.fileName);
413 1 : return { text: formatDiagnostics(name, diagnostics), data: { script: name, diagnostics } };
414 1 : }
415 :
416 1 : async function runScript(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
417 1 : const script = requireScript(state, stringArg(args, 'script'));
418 1 : const maxRows = numberArg(args, 'maxRows', DEFAULT_RUN_ROWS);
419 1 : const name = scriptDisplayName(script.fileName);
420 1 : const force = booleanArg(args, 'force');
421 1 : const conn = env.connections().connections.get(script.connection.connectionId);
422 1 : if (conn == null || conn.connectionHealth !== ConnectionHealth.ONLINE) {
423 1 : fail(`The connection ${script.connection.connectionId} of ${name} is not online, the user has to connect it`);
424 : }
425 1 : const modify: ModifySessionScripts = action => env.modifySessionScripts(sessionId, action);
426 : // What a script may do comes from its analysis, which is only current right before the run
427 1 : const beforeRun = (current: ScriptData) => {
428 : // A script with @manual is too expensive to run unfiltered, its consumers inline it with their filters
429 1 : if (current.annotations.manual && !force) {
430 1 : fail(`${name} runs only on request (@manual), it is too expensive to run in full. Read :${name} from another script of its connection that filters or aggregates it, or pass --force if the request is to run it in full.`);
431 : }
432 1 : env.onScriptRun?.(sessionId, current.scriptKey);
433 1 : };
434 1 : let table: arrow.Table | null;
435 1 : try {
436 : // Waits for the run of another script of the session, e.g. one of another turn
437 1 : table = await runSessionScriptForResult(() => env.sessionScripts(sessionId), script.scriptKey, env.executeQuery, modify, env.logger as LoggerLike, beforeRun);
438 1 : } catch (e: unknown) {
439 1 : if (e instanceof AgentCommandError) throw e;
440 1 : fail(`${name} failed: ${stringifyError(e)}`);
441 1 : }
442 1 : if (table == null) {
443 0 : return { text: `${name} ran without a result`, data: { script: name, rowCount: 0, columns: [], rows: [], offset: 0 } };
444 : }
445 1 : const slice = sliceResult(table, 0, maxRows);
446 1 : return { text: formatResultSlice(slice, name), data: { script: name, ...slice } };
447 1 : }
448 :
449 : /// The result of the last run of a script
450 1 : function lastResult(env: AgentCommandEnvironment, script: ScriptData): arrow.Table {
451 1 : const name = scriptDisplayName(script.fileName);
452 1 : const queryId = script.latestResultQueryId ?? script.latestQueryId;
453 1 : let table: arrow.Table | null = null;
454 1 : if (queryId != null) {
455 1 : for (const conn of env.connections().connections.values()) {
456 1 : const query = conn.queriesFinished.get(queryId) ?? conn.queriesActive.get(queryId);
457 1 : if (query != null) {
458 1 : if (query.error != null && query.resultTable == null) fail(`The last run of ${name} failed: ${query.error.message}`);
459 1 : table = query.resultTable;
460 1 : break;
461 1 : }
462 1 : }
463 : }
464 1 : if (table == null) fail(`${name} has no result, run it with dashql agent run ${name}`);
465 1 : return table;
466 1 : }
467 :
468 1 : function readResult(env: AgentCommandEnvironment, state: SessionScripts, args: Record<string, unknown>): AgentCommandResult {
469 1 : const script = requireScript(state, stringArg(args, 'script'));
470 1 : const page = numberArg(args, 'page', 1);
471 1 : const name = scriptDisplayName(script.fileName);
472 1 : const table = lastResult(env, script);
473 1 : const offset = (page - 1) * RESULT_PAGE_ROWS;
474 1 : if (offset >= table.numRows && table.numRows > 0) fail(`${name} has ${table.numRows} rows, page ${page} is empty`);
475 1 : const slice = sliceResult(table, offset, RESULT_PAGE_ROWS);
476 1 : return { text: formatResultSlice(slice, name), data: { script: name, page, ...slice } };
477 1 : }
478 :
479 : /// The chart of a script: its renderer and the Vega-Lite spec of its `@vis vegalite` tag, without the `$schema` and
480 : /// `data` that the compiler adds, so that the spec goes back into `vis edit` as it is
481 1 : async function readVis(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
482 1 : const script = await analyzed(env, sessionId, state, requireScript(state, stringArg(args, 'script')));
483 1 : const name = scriptDisplayName(script.fileName);
484 1 : const vis = script.annotations.visualizeQuery;
485 1 : if (vis == null) fail(`${name} has no chart, add one with dashql agent vis edit ${name}`);
486 1 : if (vis.renderer !== 'vegalite') {
487 1 : return { text: `${name} shows a ${vis.renderer} @vis tag, no Vega-Lite spec`, data: { script: name, renderer: vis.renderer, spec: null } };
488 : }
489 1 : const spec: Record<string, unknown> = { ...vis.vegaLiteSpec };
490 1 : delete spec.$schema;
491 1 : delete spec.data;
492 1 : return { text: JSON.stringify(spec, null, 2), data: { script: name, renderer: vis.renderer, spec } };
493 1 : }
494 :
495 : /// Bytes as base64
496 1 : function toBase64(bytes: Uint8Array): string {
497 1 : let binary = '';
498 1 : for (let i = 0; i < bytes.length; i += 0x8000) {
499 1 : binary += String.fromCharCode(...bytes.subarray(i, i + 0x8000));
500 : }
501 1 : return btoa(binary);
502 1 : }
503 :
504 : /// The limits of chart images, in CSS pixels and pixels per CSS pixel
505 1 : const MAX_IMAGE_SIZE = 4000;
506 1 : const MAX_IMAGE_SCALE = 4;
507 :
508 : /// Render the chart of a script as PNG: the `@vis vegalite` tag of its last query with the last result as data
509 1 : async function renderVis(env: AgentCommandEnvironment, sessionId: string, state: SessionScripts, args: Record<string, unknown>): Promise<AgentCommandResult> {
510 1 : const script = await analyzed(env, sessionId, state, requireScript(state, stringArg(args, 'script')));
511 1 : const name = scriptDisplayName(script.fileName);
512 1 : const width = numberArg(args, 'width', CHART_IMAGE_WIDTH);
513 1 : const height = numberArg(args, 'height', CHART_IMAGE_HEIGHT);
514 1 : const scale = numberArg(args, 'scale', CHART_IMAGE_SCALE);
515 1 : if (width > MAX_IMAGE_SIZE || height > MAX_IMAGE_SIZE) usage(`width and height must be at most ${MAX_IMAGE_SIZE}`);
516 1 : if (scale > MAX_IMAGE_SCALE) usage(`scale must be at most ${MAX_IMAGE_SCALE}`);
517 1 : const vis = script.annotations.visualizeQuery;
518 1 : if (vis?.renderer !== 'vegalite') fail(`${name} has no Vega-Lite chart, add one with dashql agent vis edit ${name}`);
519 1 : const table = lastResult(env, script);
520 1 : let image: ChartImage;
521 1 : try {
522 1 : image = await renderVegaLitePng(vis.vegaLiteSpec, arrowTableToRows(table), { width, height, scale });
523 1 : } catch (e: unknown) {
524 0 : fail(`Cannot render the chart of ${name}: ${stringifyError(e)}`);
525 0 : }
526 1 : const kib = Math.max(1, Math.round(image.png.length / 1024));
527 1 : return {
528 1 : text: `Rendered the chart of ${name}: ${image.width}x${image.height} pixels, ${table.numRows} rows, ${kib} KiB`,
529 1 : data: { script: name, width: image.width, height: image.height, rows: table.numRows, png: toBase64(image.png) },
530 1 : };
531 1 : }
532 :
533 : /// Run a command of the command line tool against a session
534 1 : export async function runAgentCommand(env: AgentCommandEnvironment, params: AgentCommandParams): Promise<AgentCommandResult> {
535 1 : const state = env.sessionScripts(params.session);
536 1 : if (state == null) throw new AgentCommandError('not_connected', `The session ${params.session} is not open in DashQL`);
537 1 : const args = params.args ?? {};
538 1 : switch (params.command) {
539 1 : case 'scripts':
540 1 : return listScripts(env, state);
541 1 : case 'read': {
542 1 : const script = requireScript(state, stringArg(args, 'script'));
543 1 : const text = script.scriptSession.getText();
544 1 : return { text, data: { script: scriptDisplayName(script.fileName), text } };
545 1 : }
546 1 : case 'edit':
547 1 : return editScript(env, params.session, state, args);
548 1 : case 'catalog':
549 1 : return describeCatalog(env, state, args);
550 1 : case 'lint':
551 1 : return lintScript(env, params.session, state, args);
552 1 : case 'run':
553 1 : return runScript(env, params.session, state, args);
554 1 : case 'result':
555 1 : return readResult(env, state, args);
556 1 : case 'vis.read':
557 1 : return readVis(env, params.session, state, args);
558 1 : case 'vis.edit':
559 1 : return editVis(env, params.session, state, args);
560 1 : case 'vis.image':
561 1 : return renderVis(env, params.session, state, args);
562 0 : default:
563 0 : usage(`unknown command: ${params.command}`);
564 1 : }
565 1 : }
|
