Line data Source code
1 : import * as core from '@ankoh/dashql-core';
2 : import type * as arrow from 'apache-arrow';
3 :
4 : import type { CancelQuery, QueryExecutor } from '../../connections/query_executor.js';
5 : import type { QueryExecutionArgs } from '../../connections/query_execution_args.js';
6 : import type { ClientFunctionCallDescription } from '../../connections/clfn/clfn_spec.js';
7 : import { QueryType, type PlannedStatement } from '../../query/query_execution_state.js';
8 : import type { LoggerLike } from '../../platform/logger/logger.js';
9 : import { APPLY_PLANNER_ANALYSIS, REGISTER_QUERY, REGISTER_SCRIPT_RESULT, SET_SCRIPT_EXECUTION_VISUALIZATION, scriptResultOwner, type ScriptData, type SessionScripts } from './session_scripts.js';
10 : import { resolveCompiledVisualization } from './script_compilation.js';
11 : import { scriptDisplayName } from './script_types.js';
12 : import type { ModifySessionScripts } from './session_scripts_registry.js';
13 : import { yieldToBrowser } from '../../utils/yield_to_browser.js';
14 : import type { ParquetDatabase } from '../../platform/hyperdb/hyperdb_parquet_export.js';
15 : import { isQueryCacheMiss, type QueryResultCache } from '../../query/query_result_cache.js';
16 : import { affectedScriptIndices, sessionPlanRequest } from '../inputs/session_inputs.js';
17 :
18 : /// The run that holds a session, a run of the whole session or of one of its scripts
19 : type SessionRunKind = 'session' | 'script';
20 :
21 : /// The sessions that a run holds. One run holds a session at a time, the app disables the runs of a held session.
22 1 : const activeSessions = new Map<string, SessionRunKind>();
23 1 : const activeSessionListeners = new Set<() => void>();
24 1 : const PREPARATION_YIELD_INTERVAL = 8;
25 1 : const EXECUTION_YIELD_INTERVAL = 8;
26 1 : const DEFAULT_MAX_CONCURRENT_QUERIES = 1;
27 1 : const NONE = 0xffffffff;
28 1 : const LOG_CTX = 'session_execution';
29 :
30 : /// Does a run hold a session, a run of the whole session or of one of its scripts?
31 1 : export function isSessionExecutionRunning(session: SessionScripts | string): boolean {
32 1 : return activeSessions.has(typeof session === 'string' ? session : session.sessionId);
33 1 : }
34 :
35 : /// Subscribe to the runs that start and end, returns the unsubscribe
36 1 : export function subscribeSessionExecutions(listener: () => void): () => void {
37 1 : activeSessionListeners.add(listener);
38 1 : return () => { activeSessionListeners.delete(listener); };
39 1 : }
40 :
41 : /// Hold a session for a run, false if a run holds it already. `releaseSessionExecution` ends the run.
42 1 : export function acquireSessionExecution(session: SessionScripts | string, kind: SessionRunKind = 'script'): boolean {
43 1 : const sessionId = typeof session === 'string' ? session : session.sessionId;
44 1 : if (activeSessions.has(sessionId)) return false;
45 1 : activeSessions.set(sessionId, kind);
46 1 : for (const listener of [...activeSessionListeners]) listener();
47 1 : return true;
48 1 : }
49 :
50 : /// End the run of a session
51 1 : export function releaseSessionExecution(session: SessionScripts | string): void {
52 1 : if (!activeSessions.delete(typeof session === 'string' ? session : session.sessionId)) return;
53 1 : for (const listener of [...activeSessionListeners]) listener();
54 1 : }
55 :
56 : /// Wait until no run holds a session and hold it for a run of a script, e.g. a run of the agent, which can't see
57 : /// that the session runs. Rejects if a run of the whole session holds it: that run takes every script and the user
58 : /// decides when it ends.
59 1 : export function waitForSessionExecution(sessionId: string): Promise<void> {
60 1 : return new Promise((resolve, reject) => {
61 1 : let unsubscribe = () => {};
62 1 : const attempt = () => {
63 1 : const holder = activeSessions.get(sessionId);
64 1 : if (holder === 'session') {
65 1 : unsubscribe();
66 1 : reject(new Error('The session is running'));
67 1 : } else if (holder == null) {
68 : // The first waiter holds the session once it is released, before anybody else can start a run
69 1 : unsubscribe();
70 1 : acquireSessionExecution(sessionId);
71 1 : resolve();
72 1 : }
73 1 : };
74 1 : unsubscribe = subscribeSessionExecutions(attempt);
75 1 : attempt();
76 1 : });
77 1 : }
78 :
79 : export interface SessionExecutionState {
80 : phase: 'planning' | 'running' | 'stopping' | 'succeeded' | 'failed' | 'cancelled';
81 : totalScripts: number;
82 : completedScripts: number;
83 : /// The planned statements
84 : totalStatements: number;
85 : completedStatements: number;
86 : /// The script of the statement that started last
87 : activeFileName: string | null;
88 : activeStatementIndex: number | null;
89 : activeStatementCount: number | null;
90 : error: string | null;
91 : }
92 :
93 : /// The Parquet files of the materialized script results of a run.
94 : ///
95 : /// A statement of the embedded Hyper reads the result of a script of another connection from a Parquet file
96 : /// (see docs/design/script_references.md). The files are files of the query result cache of the app, mounted at the
97 : /// mount prefix of the plan. The result of a script whose key is reusable is a new file of its key, every other result
98 : /// is a transient file. Hyper writes the files and reads the byte ranges a scan needs from disk. The run pins its
99 : /// files, so eviction and other windows leave them alone.
100 : export interface ExecutionMounts {
101 : /// The mount path of the files of the run, the mount prefix of the plan
102 : readonly prefix: string;
103 : /// The name of a new file of the result of a script, a file of its key if the key is reusable, a transient file
104 : /// otherwise
105 : fileName(key: string | null): string;
106 : /// Write a script result as Parquet file at a path below the prefix, returns the file
107 : writeTable(path: string, table: arrow.Table): Promise<Blob>;
108 : /// Keep a file that the run reads, e.g. a cached result
109 : use(path: string): Promise<void>;
110 : /// Release the files of the run. The files of keys stay in the cache, the transient files are removed. Writes
111 : /// and reads after the release fail.
112 : release(): Promise<void>;
113 : }
114 :
115 : /// The files of a run in the query result cache of the app
116 1 : export function createExecutionMounts(cache: QueryResultCache, database: ParquetDatabase): ExecutionMounts {
117 1 : const pins: Array<() => Promise<void>> = [];
118 1 : const transient: string[] = [];
119 1 : let released = false;
120 1 : const fileName = (path: string) => {
121 1 : if (!path.startsWith(`${cache.mountPath}/`) || path.slice(cache.mountPath.length + 1).includes('/')) throw new Error(`The path ${path} is no file of the run`);
122 1 : return path.slice(cache.mountPath.length + 1);
123 1 : };
124 : /// Mount the files and pin a file until the release
125 1 : const keep = async (name: string) => {
126 1 : if (released) throw new Error(`The run ended, it no longer reads ${name}`);
127 1 : await cache.mount(database);
128 1 : const unpin = await cache.pin(name);
129 1 : if (released) {
130 0 : await unpin();
131 0 : throw new Error(`The run ended, it no longer reads ${name}`);
132 0 : }
133 1 : pins.push(unpin);
134 1 : };
135 1 : return {
136 1 : prefix: cache.mountPath,
137 1 : fileName: key => {
138 1 : if (key != null) return cache.newFileName(key);
139 1 : const name = cache.newTransientFileName();
140 1 : transient.push(name);
141 1 : return name;
142 1 : },
143 1 : writeTable: async (path, table) => {
144 1 : const name = fileName(path);
145 1 : await keep(name);
146 1 : return await cache.writeTable(database, table, name);
147 1 : },
148 1 : use: async path => {
149 1 : await keep(fileName(path));
150 1 : },
151 1 : release: async () => {
152 1 : released = true;
153 1 : await Promise.all(pins.splice(0, pins.length).map(release => release()));
154 : // Every listing removes them too, e.g. those of a window that was closed during its run
155 1 : await Promise.all(transient.splice(0, transient.length).map(name => cache.remove(name).catch(() => {})));
156 1 : },
157 1 : };
158 1 : }
159 :
160 : /// The result key of a script of a plan
161 : interface ScriptResultKey {
162 : /// The key, the files of the results of the script are named after it
163 : key: string;
164 : /// The signature that the key derives from: the signature of the AST of the script and the keys of the scripts it
165 : /// references, one per line. The executor derives the key from it.
166 : signature: string;
167 : /// May a run read a file of the key instead of running the script? The script is cacheable, its connection
168 : /// caches results and so do the scripts it references.
169 : reusable: boolean;
170 : }
171 :
172 : /// The result keys of the scripts of a plan. The key of a script covers its connection, its AST and the keys of the
173 : /// scripts it references (docs/design/query_cache.md). Null for scripts that aren't planned or whose connection has
174 : /// no key.
175 1 : export async function computeScriptResultKeys(
176 1 : graph: core.buffers.execution.ExecutionGraphT,
177 1 : keyOf: (scriptIndex: number, signature: string) => Promise<string | null>,
178 1 : cachesResults: (scriptIndex: number) => boolean,
179 1 : ): Promise<Array<ScriptResultKey | null>> {
180 1 : const keys = new Map<number, Promise<ScriptResultKey | null>>();
181 1 : const resultKey = (index: number, depth: number): Promise<ScriptResultKey | null> => {
182 1 : let pending = keys.get(index);
183 1 : if (pending == null) {
184 1 : pending = (async () => {
185 1 : const script = graph.scripts[index];
186 1 : if (script == null || !script.planned || depth > graph.scripts.length) return null;
187 1 : const references = await Promise.all((script.referencedScripts ?? []).map(r => resultKey(r, depth + 1)));
188 1 : if (references.some(reference => reference == null)) return null;
189 : // The signature of a script without references is the signature of its AST
190 1 : const signature = [String(script.signature ?? ''), ...references.map(reference => reference!.key)].join('\n');
191 1 : const key = await keyOf(index, signature);
192 1 : if (key == null) return null;
193 1 : const reusable = script.cacheable && cachesResults(index) && references.every(reference => reference!.reusable);
194 1 : return { key, signature, reusable };
195 1 : })();
196 1 : keys.set(index, pending);
197 1 : }
198 1 : return pending;
199 1 : };
200 1 : return await Promise.all(graph.scripts.map((_script, index) => resultKey(index, 0)));
201 1 : }
202 :
203 : /// The mount prefix of a plan without the mounts of the query result cache. Its materializations can't run, the plan
204 : /// fails if it has any.
205 1 : export function createMountPrefix(): string {
206 1 : return `/mnt/external/dashql-run-${crypto.randomUUID()}`;
207 1 : }
208 :
209 : /// A plan of the scripts of a session
210 : interface SessionScriptsPlan {
211 : /// The execution graph, `graph.scripts[i]` is the script `i` of the request
212 : graph: core.buffers.execution.ExecutionGraphT;
213 : /// The connection ids, `node.connection` indexes them
214 : connectionIds: string[];
215 : }
216 :
217 : /// Plan the scripts of a session across their connections. Without targets, the plan covers every script.
218 1 : export function planSessionScripts(args: {
219 : instance: core.DashQL;
220 : scripts: readonly ScriptData[];
221 : files: readonly string[];
222 : mountPrefix: string;
223 : targets?: number[];
224 : materialized?: number[];
225 : /// The name of the file of the materialized result of a script below the mount prefix
226 : materializationFile?: (scriptIndex: number) => string | undefined;
227 : /// The session, its attachments are table parameters of the scripts
228 : session?: Pick<SessionScripts, 'sessionId' | 'attachments'>;
229 1 : }): SessionScriptsPlan {
230 1 : const request = sessionPlanRequest(args.scripts, args.files, args.session);
231 1 : const graph = args.instance.planSession({
232 1 : ...request,
233 1 : scripts: request.scripts.map((script, index) => ({ ...script, materializationFile: args.materializationFile?.(index) })),
234 1 : targets: args.targets,
235 1 : materialized: args.materialized,
236 1 : mountPrefix: args.mountPrefix,
237 1 : }).unpack();
238 1 : const connectionIds = request.connections.map(c => c.id);
239 1 : return { graph, connectionIds };
240 1 : }
241 :
242 : /// The errors of an invalid plan, e.g. `bar.sql: unknown script or attachment :foo`
243 1 : export function describePlanErrors(graph: core.buffers.execution.ExecutionGraphT, files: readonly string[]): string {
244 1 : const diagnostics = graph.diagnostics.filter(diagnostic => diagnostic.error);
245 1 : const messages = (diagnostics.length > 0 ? diagnostics : graph.diagnostics).map(diagnostic => {
246 1 : const message = String(diagnostic.message ?? '');
247 1 : const file = diagnostic.script !== NONE ? files[diagnostic.script] : undefined;
248 1 : return file != null && !message.startsWith(file) ? `${file}: ${message}` : message;
249 1 : });
250 1 : return messages.join('; ') || 'Invalid session execution plan';
251 1 : }
252 :
253 : /// The run of a plan: of the whole session or of a target script
254 : interface PlanRun {
255 : kind: 'session' | 'script';
256 : /// The file of the target of a run of a script
257 : target: string | null;
258 : }
259 :
260 : /// Where the statements of a plan come from, one per node
261 1 : export function describePlannedStatements(graph: core.buffers.execution.ExecutionGraphT, files: readonly string[], run: PlanRun): PlannedStatement[] {
262 1 : const statementsOf = new Map<number, number[]>();
263 1 : graph.nodes.forEach((node, index) => {
264 1 : const statements = statementsOf.get(node.script) ?? [];
265 1 : statements.push(index);
266 1 : statementsOf.set(node.script, statements);
267 1 : });
268 1 : const after = graph.nodes.map(() => [] as Array<{ node: number; reason: string }>);
269 1 : for (const edge of graph.edges) {
270 1 : const reason = core.buffers.execution.ExecutionGraphEdgeReason[edge.reason] ?? String(edge.reason);
271 1 : if (!after[edge.to]?.some(e => e.node === edge.from && e.reason === reason)) after[edge.to]?.push({ node: edge.from, reason });
272 1 : }
273 1 : const materializations = graph.materializations ?? [];
274 1 : return graph.nodes.map((node, index) => {
275 1 : const statements = statementsOf.get(node.script)!;
276 1 : return {
277 1 : run: run.kind,
278 1 : target: run.target,
279 1 : node: index,
280 1 : script: files[node.script] ?? `script ${node.script}`,
281 1 : statement: statements.indexOf(index) + 1,
282 1 : statements: statements.length,
283 1 : statementType: core.buffers.parser.StatementType[node.statementType] ?? String(node.statementType),
284 1 : result: node.scriptResult,
285 1 : cached: false,
286 1 : inlined: (node.inlinedScripts ?? []).map(script => files[script] ?? `script ${script}`),
287 1 : scans: (node.inputs ?? []).map(m => materializations[m]).filter(m => m != null).map(m => ({
288 1 : script: files[m.script] ?? `script ${m.script}`, path: String(m.path ?? ''), cached: m.node === NONE,
289 1 : })),
290 1 : materializes: node.materialization !== NONE ? String(materializations[node.materialization]?.path ?? '') : null,
291 1 : after: after[index],
292 1 : };
293 1 : });
294 1 : }
295 :
296 : /// The title of a planned statement in the query history, e.g. `bar.sql 2/3 · inlines foo.sql`
297 1 : function plannedStatementTitle(statement: PlannedStatement): string {
298 1 : const parts = [`${statement.script} ${statement.statement}/${statement.statements}`];
299 1 : if (statement.cached) parts.push('cached result');
300 1 : if (statement.inlined.length > 0) parts.push(`inlines ${statement.inlined.join(', ')}`);
301 1 : if (statement.scans.length > 0) parts.push(`scans ${statement.scans.map(scan => scan.cached ? `${scan.script} (cached)` : scan.script).join(', ')}`);
302 1 : if (statement.materializes != null) parts.push('materialized');
303 1 : return parts.join(' · ');
304 1 : }
305 :
306 : /// The description of a planned statement in the query history
307 1 : function plannedStatementDescription(statement: PlannedStatement): string {
308 1 : const parts = [`${statement.statementType} ${statement.statement} of ${statement.statements} of ${statement.script}`
309 1 : + (statement.result ? ', the result of the script' : '')
310 1 : + (statement.run === 'session' ? ', run of the session' : `, run of ${statement.target}`)];
311 1 : if (statement.cached) parts.push('served from the query result cache');
312 1 : if (statement.inlined.length > 0) parts.push(`inlines the queries of ${statement.inlined.join(', ')}`);
313 1 : for (const scan of statement.scans) parts.push(`scans the ${scan.cached ? 'cached' : 'materialized'} result of ${scan.script} at ${scan.path}`);
314 1 : if (statement.materializes != null) parts.push(`writes its result to ${statement.materializes}`);
315 1 : if (statement.after.length > 0) parts.push(`runs after ${statement.after.map(e => `node ${e.node} (${e.reason})`).join(', ')}`);
316 1 : return parts.join('; ');
317 1 : }
318 :
319 : /// The log fields of a planned statement
320 1 : export function plannedStatementLogFields(statement: PlannedStatement): Record<string, string> {
321 1 : return {
322 1 : run: statement.run === 'session' ? 'session' : `script ${statement.target}`,
323 1 : node: String(statement.node),
324 1 : script: statement.script,
325 1 : statement: `${statement.statement}/${statement.statements}`,
326 1 : statementType: statement.statementType,
327 1 : result: String(statement.result),
328 1 : cached: String(statement.cached),
329 1 : inlined: statement.inlined.join(', '),
330 1 : scans: statement.scans.map(scan => `${scan.script}${scan.cached ? ' (cached)' : ''}: ${scan.path}`).join(', '),
331 1 : materializes: statement.materializes ?? '',
332 1 : after: statement.after.map(e => `${e.node} ${e.reason}`).join(', '),
333 1 : };
334 1 : }
335 :
336 : /// The query metadata of a planned statement, the history shows where the statement comes from
337 1 : export function plannedStatementMetadata(statement: PlannedStatement): { title: string; description: string; issuer: string; planned: PlannedStatement } {
338 1 : return {
339 1 : title: plannedStatementTitle(statement),
340 1 : description: plannedStatementDescription(statement),
341 1 : issuer: statement.run === 'session' ? 'Session run' : `Run of ${statement.target}`,
342 1 : planned: statement,
343 1 : };
344 1 : }
345 :
346 : /// Log a plan before it runs: its statements and the materialized results
347 1 : function logPlan(logger: LoggerLike | undefined, sessionId: string, graph: core.buffers.execution.ExecutionGraphT, files: readonly string[], run: PlanRun, statements: PlannedStatement[]): void {
348 1 : if (logger == null) return;
349 1 : const materializations = graph.materializations ?? [];
350 1 : logger.info('Running a plan of session scripts', {
351 1 : sessionId,
352 1 : run: run.kind === 'session' ? 'session' : `script ${run.target}`,
353 1 : scripts: graph.scripts.filter(script => script.planned).length.toString(),
354 1 : nodes: graph.nodes.length.toString(),
355 1 : edges: graph.edges.length.toString(),
356 1 : materializations: materializations.map(m => `${files[m.script]}: ${m.node === NONE ? 'cached' : `node ${m.node}`} -> ${m.path}`).join(', '),
357 1 : inlined: statements.filter(s => s.inlined.length > 0).map(s => `${s.node} ${s.script} inlines ${s.inlined.join(' ')}`).join(', '),
358 1 : }, LOG_CTX);
359 1 : }
360 :
361 : /// Check a plan before anything of it runs: the visualizations of its outputs decode, including outputs that the run
362 : /// discards. Returns the visualizations of the outputs by node.
363 1 : async function checkExecutionGraph(plan: SessionScriptsPlan, cooperate: () => Promise<void>) {
364 1 : const { graph } = plan;
365 1 : const visualizations = new Map<number, ReturnType<typeof resolveCompiledVisualization>>();
366 1 : for (const [index, node] of graph.nodes.entries()) {
367 1 : if (!node.producesOutput) continue;
368 1 : const visualization = node.visualization
369 1 : ? resolveCompiledVisualization(String(node.sql ?? ''), node.visualization, node.statementId) : null;
370 1 : if (node.visualization && !visualization) {
371 1 : throw new Error('Invalid captured visualization metadata');
372 : }
373 1 : visualizations.set(index, visualization);
374 1 : if ((index + 1) % PREPARATION_YIELD_INTERVAL === 0) await cooperate();
375 1 : }
376 1 : return visualizations;
377 1 : }
378 :
379 : /// The shared state of a run: the running queries and their cancellations
380 : export interface RunQueries {
381 : /// The running queries with their connections
382 : active: Map<number, { connectionId: string; execution: Promise<unknown> }>;
383 : /// The cancellations of the queries
384 : cancellations: Map<number, Promise<void>>;
385 : }
386 :
387 : /// Run the nodes of an execution graph across their connections.
388 : ///
389 : /// A node starts when its predecessors are done, at most `maxConcurrentQueries` nodes run on a connection at once.
390 : /// The node of a materialized script writes its result as Parquet file and mounts it before the nodes that scan it
391 : /// start. The last output of every script registers as the result of the script. A failed or cancelled run cancels
392 : /// the running nodes and settles once they are done, it writes no more files.
393 1 : export async function runExecutionGraph(args: {
394 : plan: SessionScriptsPlan;
395 : scripts: readonly ScriptData[];
396 : files: readonly string[];
397 : sessionId: string;
398 : executeQuery: QueryExecutor;
399 : modifySessionScripts: ModifySessionScripts;
400 : signal: AbortSignal;
401 : queries: RunQueries;
402 : /// Cancel the running nodes, e.g. once a node failed
403 : cancelActive: () => void;
404 : /// Throws if the session or a planned script changed, `checkVersions` compares the text versions too
405 : validate: (checkVersions: boolean) => void;
406 : mounts: ExecutionMounts | null;
407 : /// Mount the attachments of a session in the embedded Hyper, absent without a compute database
408 : mountAttachments?: (sessionId: string) => Promise<void>;
409 : cooperate: () => Promise<void>;
410 : emit: (update: Partial<SessionExecutionState>, increment?: { statements?: number; scripts?: number }) => void;
411 : maxConcurrentQueries?: (connectionId: string) => number;
412 : /// Adjust the execution of the result node of a script, e.g. to cache it
413 : resultQueryArgs?: (scriptIndex: number) => Partial<QueryExecutionArgs>;
414 :
415 : /// The run of the plan, the logs and the query history show where the statements come from
416 : run: PlanRun;
417 : logger?: LoggerLike;
418 1 : }): Promise<Array<number | null>> {
419 1 : const { plan: { graph, connectionIds }, scripts, files, sessionId, signal, queries, logger } = args;
420 1 : const plannedStatements = describePlannedStatements(graph, files, args.run);
421 1 : args.emit({ phase: 'running', totalStatements: graph.nodes.length });
422 1 : const visualizations = await checkExecutionGraph(args.plan, args.cooperate);
423 : // Validate after notifications, but before the first backend call. Later edits belong to the next run.
424 1 : args.validate(true);
425 1 : const materializations = graph.materializations ?? [];
426 1 : if (materializations.length > 0 && args.mounts == null) {
427 1 : throw new Error('Scripts of other connections need the compute database to scan their results');
428 : }
429 : // The statements that scan attachments read them from the mount of the session
430 1 : if (graph.nodes.some(node => (node.attachments?.length ?? 0) > 0)) {
431 1 : if (args.mountAttachments == null) throw new Error('Scripts that read attachments need the compute database');
432 1 : await args.mountAttachments(sessionId);
433 1 : }
434 1 : logPlan(logger, sessionId, graph, files, args.run, plannedStatements);
435 : // The results that the host provides, e.g. cached results, are read in place
436 1 : for (const materialization of materializations) {
437 1 : if (materialization.node === NONE) await args.mounts!.use(String(materialization.path ?? ''));
438 : }
439 1 : const predecessorCount = graph.nodes.map(() => 0);
440 1 : const successors = graph.nodes.map(() => [] as number[]);
441 1 : const uniqueEdges = new Set<string>();
442 1 : for (const edge of graph.edges) {
443 1 : const key = `${edge.from}:${edge.to}`;
444 1 : if (uniqueEdges.has(key)) continue;
445 1 : uniqueEdges.add(key);
446 1 : predecessorCount[edge.to]++;
447 1 : successors[edge.from].push(edge.to);
448 1 : }
449 1 : const statementIndexes = scripts.map(() => [] as number[]);
450 1 : const outputNodeIndexes = scripts.map(() => -1);
451 1 : const outputQueryIds = scripts.map(() => null as number | null);
452 1 : const settledStatementCounts = scripts.map(() => 0);
453 1 : for (const [nodeIndex, node] of graph.nodes.entries()) {
454 1 : statementIndexes[node.script].push(nodeIndex);
455 1 : if (node.producesOutput) outputNodeIndexes[node.script] = nodeIndex;
456 1 : }
457 1 : const ready = graph.nodes.map((_node, index) => index).filter(index => predecessorCount[index] === 0);
458 1 : const running = connectionIds.map(() => 0);
459 1 : const maxConcurrent = connectionIds.map(connectionId => Math.max(1, args.maxConcurrentQueries?.(connectionId) ?? DEFAULT_MAX_CONCURRENT_QUERIES));
460 1 : const completions: Promise<void>[] = [];
461 1 : let runningTotal = 0;
462 1 : let settled = 0;
463 1 : let failure: unknown = null;
464 1 : let notify: (() => void) | null = null;
465 1 : const wait = () => new Promise<void>(resolve => { notify = resolve; });
466 1 : const wake = () => { const resolve = notify; notify = null; resolve?.(); };
467 1 : const launch = (nodeIndex: number) => {
468 1 : const node = graph.nodes[nodeIndex];
469 1 : const connectionId = connectionIds[node.connection];
470 1 : const scriptIndex = node.script;
471 1 : const script = scripts[scriptIndex];
472 1 : const statementIndex = statementIndexes[scriptIndex].indexOf(nodeIndex);
473 1 : const output = outputNodeIndexes[scriptIndex] === nodeIndex;
474 1 : const visualization = visualizations.get(nodeIndex) ?? null;
475 : // A materialized result must be mounted before the nodes that scan it start
476 1 : const materialization = node.materialization !== NONE ? materializations[node.materialization] : null;
477 1 : let released = false;
478 1 : const release = () => {
479 1 : if (released) return;
480 1 : released = true;
481 1 : running[node.connection]--;
482 1 : runningTotal--;
483 1 : for (const successor of successors[nodeIndex]) if (--predecessorCount[successor] === 0) ready.push(successor);
484 1 : ready.sort((a, b) => a - b);
485 1 : wake();
486 1 : };
487 1 : args.validate(false);
488 1 : args.emit({ activeFileName: files[scriptIndex],
489 1 : activeStatementIndex: statementIndex + 1, activeStatementCount: statementIndexes[scriptIndex].length });
490 1 : running[node.connection]++;
491 1 : runningTotal++;
492 1 : const planned = plannedStatements[nodeIndex];
493 1 : const extra = node.scriptResult ? args.resultQueryArgs?.(scriptIndex) ?? {} : {};
494 1 : const [queryId, execution] = args.executeQuery(connectionId, {
495 1 : query: String(node.sql ?? ''), cacheable: false, throwOnError: true, abortSignal: signal,
496 1 : clientFunctionCall: node.clientFunctionCall != null ? JSON.parse(String(node.clientFunctionCall)) as ClientFunctionCallDescription : undefined,
497 1 : analyzeResults: output, publishResults: output || materialization != null,
498 1 : onResultsReady: materialization == null ? release : undefined,
499 1 : ...extra,
500 1 : metadata: {
501 1 : queryType: QueryType.USER_PROVIDED, userProvided: true, sessionId, resultOwner: scriptResultOwner(sessionId, script.scriptKey),
502 1 : ...extra.metadata,
503 1 : ...plannedStatementMetadata(planned),
504 1 : },
505 1 : });
506 1 : logger?.info('Running a planned statement', {
507 1 : sessionId, query: queryId.toString(), connectionId, ...plannedStatementLogFields(planned),
508 1 : cacheable: String(extra.cacheable ?? false), refreshCache: String(extra.refreshCache ?? false),
509 1 : }, LOG_CTX);
510 1 : queries.active.set(queryId, { connectionId, execution });
511 1 : args.modifySessionScripts({ type: REGISTER_QUERY, value: [script.scriptKey, queryId] });
512 1 : if (output) args.modifySessionScripts({ type: SET_SCRIPT_EXECUTION_VISUALIZATION, value: [script.scriptKey, visualization] });
513 1 : const completion = execution.then(async table => {
514 1 : if (materialization != null) {
515 : // A failed or cancelled run writes no more files
516 1 : if (failure != null) throw failure;
517 1 : signal.throwIfAborted();
518 1 : if (table == null) throw new Error(`${files[scriptIndex]} has no result that other scripts could read`);
519 1 : await args.mounts!.writeTable(String(materialization.path ?? ''), table);
520 1 : logger?.info('Wrote a materialized script result', {
521 1 : sessionId, query: queryId.toString(), script: files[scriptIndex], path: String(materialization.path ?? ''), rows: table.numRows.toString(),
522 1 : }, LOG_CTX);
523 1 : }
524 : // Executors that do not implement result readiness retain the historical behavior.
525 1 : release();
526 1 : if (output) {
527 1 : await args.modifySessionScripts({ type: SET_SCRIPT_EXECUTION_VISUALIZATION, value: [script.scriptKey, visualization] });
528 1 : if (table != null) await args.modifySessionScripts({ type: REGISTER_SCRIPT_RESULT, value: [script.scriptKey, queryId] });
529 1 : outputQueryIds[scriptIndex] = queryId;
530 1 : }
531 1 : queries.active.delete(queryId);
532 1 : await queries.cancellations.get(queryId);
533 1 : queries.cancellations.delete(queryId);
534 1 : if (signal.aborted) {
535 1 : wake();
536 1 : return;
537 1 : }
538 1 : args.emit({}, { statements: 1 });
539 1 : settledStatementCounts[scriptIndex]++;
540 1 : if (settledStatementCounts[scriptIndex] === statementIndexes[scriptIndex].length) args.emit({}, { scripts: 1 });
541 1 : if (++settled % EXECUTION_YIELD_INTERVAL === 0) await args.cooperate();
542 1 : wake();
543 1 : }).catch(error => {
544 1 : release();
545 1 : queries.active.delete(queryId);
546 : // The failures of cancelled nodes are no failures of the run
547 1 : const first = failure == null && !signal.aborted;
548 1 : logger?.[first ? 'warn' : 'info'](first ? 'A planned statement failed' : 'A planned statement stopped', {
549 1 : sessionId, query: queryId.toString(), connectionId, ...plannedStatementLogFields(planned),
550 1 : error: error instanceof Error ? error.message : String(error),
551 1 : }, LOG_CTX);
552 1 : if (first) { failure = error; args.cancelActive(); }
553 1 : wake();
554 1 : throw error;
555 1 : });
556 1 : completion.catch(() => {});
557 1 : completions.push(completion);
558 1 : };
559 1 : try {
560 1 : while ((ready.length > 0 || runningTotal > 0) && failure == null && !signal.aborted) {
561 : // The ready nodes of the connections that have room, in node order
562 1 : for (let i = 0; i < ready.length && failure == null && !signal.aborted;) {
563 1 : const node = graph.nodes[ready[i]];
564 1 : if (running[node.connection] >= maxConcurrent[node.connection]) { ++i; continue; }
565 1 : launch(ready.splice(i, 1)[0]);
566 1 : }
567 1 : if (runningTotal === 0) break;
568 1 : await wait();
569 1 : }
570 1 : } catch (error) {
571 : // E.g. the session changed, the nodes that run already are cancelled
572 1 : if (failure == null) failure = error;
573 1 : }
574 1 : if (failure != null || signal.aborted) {
575 : // The running nodes are done before the caller releases the files of the run
576 1 : if (runningTotal > 0) logger?.info('Cancelling the running statements of the plan', { sessionId, running: runningTotal.toString() }, LOG_CTX);
577 1 : args.cancelActive();
578 1 : await Promise.allSettled(completions);
579 1 : throw failure ?? signal.reason ?? new DOMException('Aborted', 'AbortError');
580 1 : }
581 1 : if (ready.length > 0) throw new Error('Execution plan has a cycle');
582 1 : await Promise.all(completions);
583 1 : for (const [scriptIndex, script] of scripts.entries()) {
584 1 : const outputQueryId = outputQueryIds[scriptIndex];
585 1 : const outputNodeIndex = outputNodeIndexes[scriptIndex];
586 1 : const statements = statementIndexes[scriptIndex];
587 1 : if (outputQueryId !== null && outputNodeIndex !== statements[statements.length - 1]) {
588 1 : await args.modifySessionScripts({ type: REGISTER_QUERY, value: [script.scriptKey, outputQueryId] });
589 1 : await args.modifySessionScripts({ type: SET_SCRIPT_EXECUTION_VISUALIZATION,
590 1 : value: [script.scriptKey, visualizations.get(outputNodeIndex) ?? null] });
591 1 : }
592 1 : }
593 1 : return outputQueryIds;
594 1 : }
595 :
596 : /// Validate a plan against the live session: the session keeps its scripts, their order and their connections, and
597 : /// the planned scripts keep the text of the plan. The catalog may move on, the plan holds the statements as written and
598 : /// the catalog follows them while they run.
599 1 : export function sessionPlanValidator(args: {
600 : graph: core.buffers.execution.ExecutionGraphT;
601 : session: SessionScripts;
602 : getSession: () => SessionScripts | null;
603 : scripts: readonly ScriptData[];
604 : files: readonly string[];
605 1 : }): (checkVersions: boolean) => void {
606 1 : const { graph, session, scripts, files } = args;
607 1 : return (checkVersions: boolean) => {
608 1 : const live = args.getSession();
609 1 : if (!live || live.sessionId !== session.sessionId || live.instance !== session.instance
610 1 : || live.scriptOrder.join('\0') !== files.join('\0') || files.some(file => live.scriptRefs[file] == null)) {
611 1 : throw new Error('Session changed during execution');
612 : }
613 1 : scripts.forEach((script, index) => {
614 1 : if (live.scriptRefs[files[index]].scriptId !== script.scriptKey
615 1 : || live.scripts[script.scriptKey]?.scriptSession !== script.scriptSession) {
616 0 : throw new Error(`Session script changed during execution: ${files[index]}`);
617 : }
618 1 : });
619 1 : if (graph.scripts.length !== scripts.length) {
620 0 : throw new Error('Session changed during execution');
621 : }
622 1 : if (!checkVersions) return;
623 1 : scripts.forEach((script, index) => {
624 1 : const metadata = graph.scripts[index];
625 1 : if (!metadata.planned) return;
626 1 : const scriptSession = script.scriptSession;
627 1 : if (metadata.scriptId !== scriptSession.getCatalogEntryId()
628 1 : || metadata.textVersion !== scriptSession.getTextVersion()
629 1 : || metadata.analyzedTextVersion !== metadata.textVersion) {
630 1 : throw new Error(`Stale execution plan: ${files[index]}`);
631 : }
632 1 : });
633 1 : };
634 1 : }
635 :
636 : /// Analyze every script of a session and publish the declarations into the catalogs of their connections. Returns
637 : /// the editor updates of the analysis at the final catalog revisions.
638 1 : export async function prepareSessionScripts(
639 1 : scripts: readonly ScriptData[],
640 1 : cooperate: () => Promise<void>,
641 1 : ): Promise<Array<{ scriptKey: number; editorUpdate: core.buffers.editor.EditorUpdateT }>> {
642 1 : const analyze = (script: ScriptData) => {
643 1 : const update = script.scriptSession.analyze();
644 1 : if (update.status !== core.buffers.editor.EditorUpdateStatus.OK || !update.analysisAvailable) {
645 1 : throw new Error(String(update.statusMessage || `Failed to analyze ${script.fileName}`));
646 : }
647 1 : return update;
648 1 : };
649 1 : for (const [index, script] of scripts.entries()) {
650 1 : analyze(script);
651 1 : script.scriptSession.loadIntoCatalog();
652 1 : if ((index + 1) % PREPARATION_YIELD_INTERVAL === 0) await cooperate();
653 1 : }
654 : // Publishing changed the catalogs, refresh the analysis at the current catalog revisions
655 1 : const updates: Array<{ scriptKey: number; editorUpdate: core.buffers.editor.EditorUpdateT }> = [];
656 1 : for (const [index, script] of scripts.entries()) {
657 1 : updates.push({ scriptKey: script.scriptKey, editorUpdate: analyze(script) });
658 1 : if ((index + 1) % PREPARATION_YIELD_INTERVAL === 0) await cooperate();
659 1 : }
660 1 : return updates;
661 1 : }
662 :
663 : /// A script that can't be planned, e.g. one that references an unknown script
664 1 : export class ScriptPlanError extends Error {
665 1 : constructor(message: string) {
666 1 : super(message);
667 1 : this.name = 'ScriptPlanError';
668 1 : }
669 1 : }
670 :
671 : /// Run a plan of the scripts of a session: analyze the scripts, plan the targets with what they need, every script
672 : /// without targets, serve the cached results and run the rest.
673 : ///
674 : /// A script with a reusable cached result doesn't run if no script runs that it follows in the plan: its result is
675 : /// served from the cache, a script that scans it reads the cached file. A refresh reads no cached results. A script
676 : /// that runs replaces the result of its reusable key. Releases the mounts and the cached files once the run settled.
677 : /// Returns the queries of the outputs of the scripts.
678 1 : export async function runSessionPlan(args: {
679 : session: SessionScripts;
680 : getSession: () => SessionScripts | null;
681 : /// The scripts of the plan with what they need, every script if absent
682 : targets?: number[];
683 : /// Plan the scripts that the change of these inputs affects instead: the scripts that use them and the scripts
684 : /// downstream of them (docs/design/session_inputs.md)
685 : inputs?: readonly string[];
686 : /// The run of the plan, a run of a script needs a statement of its target
687 : run: PlanRun;
688 : executeQuery: QueryExecutor;
689 : modifySessionScripts: ModifySessionScripts;
690 : signal: AbortSignal;
691 : queries: RunQueries;
692 : /// Cancel the running queries, e.g. once a statement failed
693 : cancelActive: () => void;
694 : cooperate: () => Promise<void>;
695 : emit: (update: Partial<SessionExecutionState>, increment?: { statements?: number; scripts?: number }) => void;
696 : /// The mounts of the materialized script results of the run, null without a compute database
697 : mounts: ExecutionMounts | null;
698 : /// Mount the attachments of a session in the embedded Hyper, absent without a compute database
699 : mountAttachments?: (sessionId: string) => Promise<void>;
700 : /// The result key of a script on a connection, names the files of the materialized results
701 : resultKey: (connectionId: string, signature: string) => Promise<string | null>;
702 : /// Does a connection cache query results? The run reuses and writes the results of the scripts with reusable keys.
703 : cachesResults: (connectionId: string) => boolean;
704 : /// The query result cache of the app, null if results aren't cached
705 : cache: QueryResultCache | null;
706 : /// Read no cached results: the scripts run, the results with reusable keys become the newest files of their keys
707 : refresh: boolean;
708 : maxConcurrentQueries?: (connectionId: string) => number;
709 : logger?: LoggerLike;
710 1 : }): Promise<Array<number | null>> {
711 1 : const { session, signal, queries, cooperate, emit, mounts, cache, run } = args;
712 1 : const { sessionId, instance } = session;
713 : // The cached files that the run serves or scans, pinned from the lookup on
714 1 : const retained: Array<() => Promise<void>> = [];
715 1 : try {
716 1 : const files = [...session.scriptOrder];
717 1 : emit({ totalScripts: files.length });
718 : // Let React commit the running state and the browser paint its spinner before synchronous
719 : // analysis and planning work starts.
720 1 : await cooperate();
721 1 : const scripts = files.map(file => {
722 1 : const script = session.scripts[session.scriptRefs[file].scriptId];
723 1 : if (!script) throw new Error(`Missing script: ${file}`);
724 1 : return script;
725 1 : });
726 1 : const updates = await prepareSessionScripts(scripts, cooperate);
727 1 : await cooperate();
728 : // A change of inputs runs the scripts that it affects, with what they read
729 1 : let targets = args.targets;
730 1 : if (args.inputs != null) {
731 1 : const dataflow = instance.deriveSessionDataflow(sessionPlanRequest(scripts, files, session)).unpack() as core.buffers.execution.SessionDataflowT;
732 1 : targets = affectedScriptIndices(dataflow, args.inputs);
733 1 : emit({ totalScripts: targets.length });
734 1 : if (targets.length === 0) {
735 1 : await args.modifySessionScripts({ type: APPLY_PLANNER_ANALYSIS, value: updates });
736 1 : return files.map(() => null);
737 1 : }
738 1 : }
739 1 : const mountPrefix = mounts?.prefix ?? createMountPrefix();
740 1 : const planWith = (provided: number[] = [], names: Map<number, string> = new Map()) => {
741 : // A catalog update may have landed since the preparation, the planner needs the current analysis
742 1 : for (const script of scripts) script.scriptSession.analyze();
743 1 : const plan = planSessionScripts({
744 1 : instance, scripts, files, mountPrefix, targets, materialized: provided, materializationFile: index => names.get(index), session,
745 1 : });
746 1 : if (!plan.graph.valid) throw new ScriptPlanError(describePlanErrors(plan.graph, files));
747 1 : return plan;
748 1 : };
749 : // The plan of the targets checks the session and gives the result keys
750 1 : let plan = planWith();
751 1 : if (run.kind === 'script' && targets?.every(target => !plan.graph.nodes.some(node => node.script === target))) {
752 1 : throw new ScriptPlanError(`${run.target} holds no statement`);
753 : }
754 1 : const keys = await computeScriptResultKeys(plan.graph,
755 1 : (index, signature) => args.resultKey(scripts[index].connection.connectionId, signature),
756 1 : index => args.cachesResults(scripts[index].connection.connectionId));
757 1 : const visualizations = await checkExecutionGraph(plan, cooperate);
758 1 : const validate = sessionPlanValidator({ graph: plan.graph, session, getSession: args.getSession, scripts, files });
759 1 : validate(true);
760 1 : await args.modifySessionScripts({ type: APPLY_PLANNER_ANALYSIS, value: updates });
761 1 : signal.throwIfAborted();
762 1 : validate(true);
763 : // The scripts whose cached results are served, with their files
764 1 : const cachedFiles = new Map<number, string>();
765 1 : if (!args.refresh && cache != null) {
766 1 : const statements = describePlannedStatements(plan.graph, files, run);
767 : // The scripts of the nodes that the nodes of a script follow, edges go from earlier to later scripts
768 1 : const follows = scripts.map(() => new Set<number>());
769 1 : for (const edge of plan.graph.edges) {
770 1 : const from = plan.graph.nodes[edge.from].script, to = plan.graph.nodes[edge.to].script;
771 1 : if (from !== to) follows[to].add(from);
772 1 : }
773 1 : const runs = scripts.map(() => true);
774 1 : for (const [index, script] of scripts.entries()) {
775 1 : const key = keys[index];
776 1 : const resultNodeIndex = plan.graph.nodes.findIndex(node => node.script === index && node.scriptResult);
777 1 : if (!key?.reusable || resultNodeIndex < 0 || [...follows[index]].some(earlier => runs[earlier])) continue;
778 1 : const file = await cache.retain(key.key).catch(() => null);
779 1 : signal.throwIfAborted();
780 1 : if (file == null) continue;
781 1 : retained.push(file.release);
782 1 : const resultNode = plan.graph.nodes[resultNodeIndex];
783 1 : const statement = { ...statements[resultNodeIndex], cached: true };
784 1 : const visualization = visualizations.get(resultNodeIndex) ?? null;
785 1 : const connectionId = script.connection.connectionId;
786 1 : const [queryId, execution] = args.executeQuery(connectionId, {
787 1 : query: String(resultNode.sql ?? ''), analyzeResults: true, publishResults: true, throwOnError: true, abortSignal: signal,
788 1 : cacheable: true, cacheOnly: true, cacheSignature: key.signature,
789 1 : metadata: {
790 1 : queryType: QueryType.USER_PROVIDED, userProvided: true, sessionId, resultOwner: scriptResultOwner(sessionId, script.scriptKey),
791 1 : ...plannedStatementMetadata(statement),
792 1 : },
793 1 : });
794 1 : args.logger?.info('Serving the cached result of a script', {
795 1 : sessionId, query: queryId.toString(), key: key.key, ...plannedStatementLogFields(statement),
796 1 : }, LOG_CTX);
797 1 : queries.active.set(queryId, { connectionId, execution });
798 1 : const served = await execution.then(table => ({ table }), async (error: unknown) => {
799 1 : if (!isQueryCacheMiss(error)) {
800 : // The script shows the error of the query
801 0 : await args.modifySessionScripts({ type: REGISTER_QUERY, value: [script.scriptKey, queryId] });
802 0 : throw error;
803 0 : }
804 1 : args.logger?.info('The cached result of a script is gone, running it', { sessionId, query: queryId.toString(), script: files[index], key: key.key }, LOG_CTX);
805 1 : return null;
806 1 : }).finally(() => queries.active.delete(queryId));
807 1 : if (served == null) continue;
808 1 : runs[index] = false;
809 1 : cachedFiles.set(index, file.entry.name);
810 1 : await args.modifySessionScripts({ type: REGISTER_QUERY, value: [script.scriptKey, queryId] });
811 1 : await args.modifySessionScripts({ type: SET_SCRIPT_EXECUTION_VISUALIZATION, value: [script.scriptKey, visualization] });
812 1 : if (served.table != null) await args.modifySessionScripts({ type: REGISTER_SCRIPT_RESULT, value: [script.scriptKey, queryId] });
813 1 : await cooperate();
814 1 : }
815 1 : }
816 : // Scripts without statements are done right away, so are the scripts with cached results
817 1 : const counted = targets ?? scripts.map((_script, index) => index);
818 : // A plan of targets leaves out the targets with cached results
819 1 : if (targets != null) targets = targets.filter(target => !cachedFiles.has(target));
820 1 : if (targets?.length === 0) {
821 1 : emit({}, { scripts: counted.length });
822 1 : return files.map(() => null);
823 1 : }
824 : // The plan of the scripts that run, it scans the cached results. A result that the run materializes is a new
825 : // file of its key if the key is reusable, a transient file otherwise.
826 1 : if (cachedFiles.size > 0 || (mounts != null && plan.graph.materializations.length > 0)) {
827 1 : const names = new Map(cachedFiles);
828 1 : for (const m of mounts != null ? plan.graph.materializations : []) {
829 1 : if (!names.has(m.script)) names.set(m.script, mounts!.fileName(keys[m.script]?.reusable ? keys[m.script]!.key : null));
830 : }
831 1 : plan = planWith([...cachedFiles.keys()], names);
832 1 : }
833 : // A materialized result is the file of its key already
834 1 : const materialized = new Set(plan.graph.materializations.map(m => m.script));
835 1 : const graph = plan.graph;
836 1 : const done = counted.filter(index => !graph.nodes.some(node => node.script === index)).length;
837 1 : if (done > 0) emit({}, { scripts: done });
838 1 : const outputs = await runExecutionGraph({
839 1 : plan, scripts, files, sessionId,
840 1 : executeQuery: args.executeQuery, modifySessionScripts: args.modifySessionScripts,
841 1 : signal, queries, cancelActive: args.cancelActive,
842 1 : validate: sessionPlanValidator({ graph, session, getSession: args.getSession, scripts, files }),
843 1 : mounts, mountAttachments: args.mountAttachments, cooperate, emit,
844 1 : maxConcurrentQueries: args.maxConcurrentQueries,
845 1 : run, logger: args.logger,
846 : // The cached results were read already, a script that runs replaces the result of its key
847 1 : resultQueryArgs: index => keys[index]?.reusable && !materialized.has(index)
848 1 : ? { cacheable: true, refreshCache: true, cacheSignature: keys[index]!.signature }
849 1 : : {},
850 1 : });
851 1 : return outputs;
852 1 : } finally {
853 1 : await mounts?.release();
854 1 : await Promise.all(retained.map(release => release().catch(() => {})));
855 1 : }
856 1 : }
857 :
858 : /// Run every script of a session, or the scripts that a change of inputs affects, and report the state of the run
859 1 : export async function runSessionExecution(args: {
860 : session: SessionScripts;
861 : getSession: () => SessionScripts | null;
862 : executeQuery: QueryExecutor;
863 : cancelQuery: CancelQuery;
864 : modifySessionScripts: ModifySessionScripts;
865 : signal: AbortSignal;
866 : onState: (state: SessionExecutionState) => void;
867 : /// The mounts of the materialized script results of the run, null without a compute database
868 : mounts?: ExecutionMounts | null;
869 : /// Mount the attachments of a session in the embedded Hyper, absent without a compute database
870 : mountAttachments?: (sessionId: string) => Promise<void>;
871 : /// The result key of a script on a connection, names the files of the materialized results
872 : resultKey?: (connectionId: string, signature: string) => Promise<string | null>;
873 : /// Does a connection cache query results? The run reuses and writes the results of the scripts with reusable keys.
874 : cachesResults?: (connectionId: string) => boolean;
875 : /// The query result cache of the app, null if results aren't cached
876 : cache?: QueryResultCache | null;
877 : /// Read no cached results: every script runs, the results with reusable keys become the newest files of their keys
878 : refresh?: boolean;
879 : /// Run only the scripts that the change of these inputs affects
880 : inputs?: readonly string[];
881 : /// Allows tests to control cooperative scheduling without relying on timer timing.
882 : yieldToBrowser?: () => Promise<void>;
883 : /// Maximum concurrently running backend queries on a connection.
884 : maxConcurrentQueries?: (connectionId: string) => number;
885 : /// Logs the plan and its statements
886 : logger?: LoggerLike;
887 1 : }): Promise<void> {
888 1 : const { session, signal } = args;
889 1 : if (!acquireSessionExecution(session, 'session')) return;
890 1 : let state: SessionExecutionState = {
891 1 : phase: 'planning', totalScripts: 0, completedScripts: 0,
892 1 : totalStatements: 0, completedStatements: 0, activeFileName: null,
893 1 : activeStatementIndex: null, activeStatementCount: null, error: null,
894 1 : };
895 1 : const emit = (update: Partial<SessionExecutionState>, increment: { statements?: number; scripts?: number } = {}) => {
896 1 : state = {
897 1 : ...state, ...update,
898 1 : completedStatements: state.completedStatements + (increment.statements ?? 0),
899 1 : completedScripts: state.completedScripts + (increment.scripts ?? 0),
900 1 : };
901 1 : args.onState(state);
902 1 : };
903 1 : const queries: RunQueries = { active: new Map(), cancellations: new Map() };
904 1 : const cooperate = async () => {
905 1 : signal.throwIfAborted();
906 1 : await (args.yieldToBrowser ?? yieldToBrowser)();
907 1 : signal.throwIfAborted();
908 1 : };
909 1 : const cancelActive = () => {
910 1 : for (const [queryId, { connectionId }] of queries.active) {
911 1 : if (queries.cancellations.has(queryId)) continue;
912 : // Catch immediately: cancellation can reject before execution has drained.
913 1 : try {
914 1 : queries.cancellations.set(queryId, args.cancelQuery(connectionId, queryId).catch(() => {}));
915 1 : } catch {
916 1 : queries.cancellations.set(queryId, Promise.resolve());
917 1 : }
918 1 : }
919 1 : };
920 1 : const stop = () => {
921 1 : emit({ phase: 'stopping' });
922 1 : cancelActive();
923 1 : };
924 1 : signal.addEventListener('abort', stop, { once: true });
925 1 : try {
926 1 : await runSessionPlan({
927 1 : session, getSession: args.getSession, inputs: args.inputs, run: { kind: 'session', target: null },
928 1 : executeQuery: args.executeQuery, modifySessionScripts: args.modifySessionScripts,
929 1 : signal, queries, cancelActive, cooperate, emit,
930 1 : mounts: args.mounts ?? null, mountAttachments: args.mountAttachments,
931 1 : resultKey: args.resultKey ?? (async () => null), cachesResults: args.cachesResults ?? (() => false),
932 1 : cache: args.cache ?? null, refresh: args.refresh ?? false,
933 1 : maxConcurrentQueries: args.maxConcurrentQueries, logger: args.logger,
934 1 : });
935 1 : signal.throwIfAborted();
936 1 : signal.removeEventListener('abort', stop);
937 1 : emit({ phase: 'succeeded', activeFileName: null,
938 1 : activeStatementIndex: null, activeStatementCount: null });
939 1 : } catch (error) {
940 1 : cancelActive();
941 1 : await Promise.allSettled([...queries.active.values()].map(({ execution }) => execution));
942 1 : await Promise.allSettled(queries.cancellations.values());
943 1 : const cancelled = signal.aborted || (error instanceof Error && error.name === 'AbortError');
944 1 : signal.removeEventListener('abort', stop);
945 1 : emit({ phase: cancelled ? 'cancelled' : 'failed',
946 1 : error: cancelled ? null : error instanceof Error ? error.message : String(error) });
947 1 : } finally {
948 1 : signal.removeEventListener('abort', stop);
949 1 : releaseSessionExecution(session);
950 1 : }
951 1 : }
|
