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