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