Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import * as React from 'react';
3 :
4 : import type * as core from '@ankoh/dashql-core';
5 :
6 : import { useConnectionDispatch, useConnectionRegistryStore, useDynamicConnectionDispatch } from './connection_registry.js';
7 : import { APPLY_CATALOG_CHANGES } from './connection_state.js';
8 : import {
9 : createQueryResponseStreamMetrics,
10 : QUERY_CACHE_RECORDED,
11 : QUERY_PROCESSING_RESULTS,
12 : QUERY_PROCESSED_RESULTS,
13 : QUERY_PROGRESS_UPDATED,
14 : QUERY_RECEIVED_ALL_BATCHES,
15 : QUERY_RECEIVED_BATCH,
16 : QueryExecutionResponseStream,
17 : QueryExecutionTracker,
18 : } from '../query/query_execution_state.js';
19 : import { useComputationDispatch } from '../compute/computation_registry.js';
20 : import { analyzeTable } from '../compute/analyze_table.js';
21 : import { DELETE_COMPUTATION } from '../compute/computation_state.js';
22 : import { ResultComputations } from './result_computations.js';
23 : import { useComputeDatabase } from '../compute/compute_connection_provider.js';
24 : import { QueryCacheMissError, type QueryResultCacheEntry } from '../query/query_result_cache.js';
25 : import { getConnectionParamsFromStateDetails, createConnectionParamsSignature } from './connection_params.js';
26 : import { ConnectionStateDetailsVariant } from './connection_state_details.js';
27 : import { computeQueryResultCacheKey } from '../query/query_result_cache_key.js';
28 : import { useLogger } from '../platform/logger/logger_provider.js';
29 : import { createTrace, type TraceContext } from '../platform/logger/trace_context.js';
30 : import { QueryExecutionArgs } from './query_execution_args.js';
31 : import { LoggableException, stringifyError, type TracedLogger } from '../platform/logger/logger.js';
32 : import type { LogRecord } from '../platform/logger/log_buffer.js';
33 : import { allocateQueryId, executeTrackedQuery } from '../query/tracked_query_execution.js';
34 : import { consumeQueryResponseStream, executeConnectionQuery } from './query_execution.js';
35 : import { yieldToBrowser } from '../utils/yield_to_browser.js';
36 : import { supportsParquet, type ParquetDatabase } from '../platform/hyperdb/hyperdb_parquet_export.js';
37 : import { useQueryResultCache } from '../storage/query_result_cache_provider.js';
38 : import { QueryProgressModels, type EngineProgressPayload } from '../query/query_progress_model.js';
39 : import { createProgressLayoutConfig } from '../views/query_profile/profile_layout_config.js';
40 : import { useDashQLCoreSetup, type DashQLSetupFn } from '../runtime/core_provider.js';
41 :
42 1 : const LOG_CTX = 'query_executor';
43 :
44 : /// Query failures are logged as warnings.
45 : /// The failed run already shows the error inline, error toasts are reserved for failures of the app itself.
46 1 : export function logQueryFailure(logger: TracedLogger, error: unknown, queryId: number, sessionId: string): void {
47 1 : if (error instanceof LoggableException) {
48 1 : logger.warn(error.message, {
49 1 : ...error.keyValues,
50 1 : queryId: queryId.toString(),
51 1 : sessionId,
52 1 : }, error.target);
53 : } else {
54 1 : logger.warn('Query failed with unknown error', {
55 1 : queryId: queryId.toString(),
56 1 : sessionId,
57 1 : raw: stringifyError(error),
58 1 : }, LOG_CTX);
59 : }
60 1 : }
61 :
62 : /// Compute the file-based cache key for a query against a connection, or null when the connection has
63 : /// no recoverable params/signature (e.g. before setup completes). This is the same derivation the
64 : /// executor uses on its cacheable path. Never throws — a failure to derive the key is reported as
65 : /// null (treat as "not cacheable / a miss").
66 1 : export async function computeQueryCacheKeyForConnection(
67 1 : details: ConnectionStateDetailsVariant,
68 1 : cacheSignature: string,
69 1 : ): Promise<string | null> {
70 1 : try {
71 1 : const params = getConnectionParamsFromStateDetails(details);
72 1 : const sig = params ? createConnectionParamsSignature(params) : null;
73 1 : if (sig == null) {
74 1 : return null;
75 : }
76 1 : return await computeQueryResultCacheKey(sig, cacheSignature);
77 1 : } catch {
78 1 : return null;
79 1 : }
80 1 : }
81 :
82 : /// The query executor function
83 : export type QueryExecutor = (connectionId: string, args: QueryExecutionArgs) => [number, Promise<arrow.Table | null>];
84 : export type CancelQuery = (connectionId: string, queryId: number) => Promise<void>;
85 : interface QueryExecutionRuntime {
86 : /// Aborted by cancel() and by the abort signal of the caller, both stop the query at the backend
87 : cancellation: AbortController;
88 : resultStream: QueryExecutionResponseStream | null;
89 : /// The cancellation of the result stream at the backend, sent once
90 : backendCancellation: Promise<void> | null;
91 : execution: Promise<arrow.Table | null> | null;
92 : /// The views that show the profile of the query
93 : profileViewers: number;
94 : }
95 : /// Show the profile of a running query, returns the function that stops showing it
96 : export type WatchQueryProfile = (queryId: number) => () => void;
97 : /// The React context to resolve the active query executor
98 : export interface QueryExecutorContextValue {
99 : execute: QueryExecutor;
100 : cancel: CancelQuery;
101 : watchProfile: WatchQueryProfile;
102 : }
103 1 : export const EXECUTOR_CTX = React.createContext<QueryExecutorContextValue | null>(null);
104 : /// The hook to resolve the query executor
105 1 : export const useQueryExecutor = () => React.useContext(EXECUTOR_CTX)!.execute;
106 1 : export const useCancelQuery = () => React.useContext(EXECUTOR_CTX)!.cancel;
107 : /// Show the profile of a query while the component is mounted: the connector then fetches the plan and its
108 : /// statistics (Trino). Without executor, e.g. in tests, nothing is fetched.
109 1 : export function useQueryProfileDetail(queryId: number | null, enabled = true) {
110 1 : const watchProfile = React.useContext(EXECUTOR_CTX)?.watchProfile;
111 1 : React.useEffect(() => {
112 1 : if (!enabled || queryId == null || watchProfile == null) return;
113 1 : return watchProfile(queryId);
114 1 : }, [watchProfile, queryId, enabled]);
115 1 : }
116 : /// Use the query state
117 1 : export function useQueryState(sessionId: string | null, queryId: number | null) {
118 : // Subscribe to the single query instead of the registry.
119 : // The registry changes with every lifecycle event of every query, including the internal queries of computations.
120 1 : const store = useConnectionRegistryStore();
121 1 : const getQueryState = React.useCallback(() => {
122 1 : if (queryId == null || sessionId == null) return null;
123 : // Query ids are unique across connections. A script that moved to another connection keeps
124 : // the result of its last query on the connection it left.
125 1 : for (const database of store.current.connections.values()) {
126 1 : const query = database.queriesActive.get(queryId) ?? database.queriesFinished.get(queryId);
127 1 : if (query != null) return query;
128 1 : }
129 1 : return null;
130 1 : }, [store, sessionId, queryId]);
131 1 : return React.useSyncExternalStore(store.subscribe, getQueryState);
132 1 : }
133 :
134 1 : function createConnectionQueryExecutionTracker(connectionId: string, dispatch: ReturnType<typeof useDynamicConnectionDispatch>[1]): QueryExecutionTracker {
135 1 : return {
136 1 : dispatch: action => dispatch(connectionId, action),
137 1 : };
138 1 : }
139 :
140 1 : export function QueryExecutorProvider(props: { children?: React.ReactElement }) {
141 1 : const logger = useLogger();
142 :
143 : // The connection registry changes frequently.
144 : // This executor resolves everything ad-hoc through the registry store and never re-renders for registry changes.
145 1 : const connStore = useConnectionRegistryStore();
146 1 : const connDispatch = useConnectionDispatch();
147 :
148 1 : const computeDispatch = useComputationDispatch();
149 1 : const computeDb = useComputeDatabase();
150 1 : const queryResultCache = useQueryResultCache();
151 1 : const queryRuntimes = React.useRef(new Map<number, QueryExecutionRuntime>());
152 1 : const resultComputations = React.useRef(new ResultComputations());
153 :
154 : // The progress models of the running queries, they need the core. The core is requested with the first
155 : // structured progress, until it is ready the progress has no structure.
156 1 : const coreSetup = useDashQLCoreSetup() as DashQLSetupFn | null;
157 1 : const progressModels = React.useRef(new QueryProgressModels(createProgressLayoutConfig()));
158 1 : const progressCore = React.useRef<core.DashQL | null>(null);
159 1 : const progressCoreRequested = React.useRef(false);
160 1 : const applyEngineProgress = React.useCallback((queryId: number, payload: EngineProgressPayload) => {
161 0 : if (progressCore.current == null && !progressCoreRequested.current && coreSetup != null) {
162 0 : progressCoreRequested.current = true;
163 0 : coreSetup('query_progress').then(
164 0 : instance => { progressCore.current = instance; },
165 0 : () => { /* Queries run without structured progress */ },
166 0 : );
167 0 : }
168 0 : return progressModels.current.apply(progressCore.current, queryId, payload);
169 1 : }, [coreSetup]);
170 :
171 : // Execute a query with pre-allocated query id
172 1 : const executeImpl = React.useCallback(async (connectionId: string, args: QueryExecutionArgs, queryId: number, runtime: QueryExecutionRuntime, trace: TraceContext): Promise<arrow.Table | null> => {
173 1 : const traced = logger.withTrace(trace);
174 1 : if (!computeDb) {
175 1 : throw new Error(`Compute database is not yet ready`);
176 : }
177 1 : const conn = connStore.current.connections.get(connectionId);
178 1 : if (!conn) {
179 1 : traced.error("Connection not configured", { connectionId, "query": queryId.toString() }, LOG_CTX);
180 1 : throw new Error(`Couldn't find a connection with id ${connectionId}`);
181 1 : }
182 : // Results are cached per session, queries without a session are never cached
183 1 : const querySessionId = args.metadata.sessionId ?? null;
184 : // Log contexts carry strings
185 1 : const sessionId = querySessionId ?? '';
186 1 : const queryTracker = createConnectionQueryExecutionTracker(connectionId, connDispatch);
187 1 : const execution = await executeTrackedQuery({
188 1 : query: args.query,
189 1 : metadata: args.metadata,
190 1 : tracker: queryTracker,
191 1 : queryId,
192 1 : trace,
193 1 : cancellation: runtime.cancellation,
194 1 : autoStart: false,
195 1 : errorTarget: LOG_CTX,
196 1 : execute: async tracked => {
197 1 : if (args.failure != null) throw args.failure;
198 1 : const initialState = tracked.initialState;
199 1 : let queryExecutionStartedAt: number | null = null;
200 1 : let queryExecutionReported = false;
201 1 : const startQueryExecution = () => {
202 1 : queryExecutionStartedAt = performance.now();
203 1 : };
204 1 : const reportQueryExecution = () => {
205 1 : if (queryExecutionStartedAt == null || queryExecutionReported) return;
206 1 : queryExecutionReported = true;
207 1 : try {
208 1 : args.onTiming?.('Results', performance.now() - queryExecutionStartedAt);
209 1 : } catch {
210 : // Timing observers must not interrupt query execution.
211 1 : }
212 1 : };
213 1 : const measurePostprocessing = async (process: () => Promise<void>) => {
214 1 : const startedAt = performance.now();
215 1 : try {
216 1 : await process();
217 1 : } finally {
218 1 : try {
219 1 : args.onTiming?.('Postprocess', performance.now() - startedAt);
220 1 : } catch {
221 : // Timing observers must not interrupt query execution.
222 0 : }
223 1 : }
224 1 : };
225 1 : let resultsReady = false;
226 1 : const reportResultsReady = () => {
227 1 : if (resultsReady) return;
228 1 : resultsReady = true;
229 1 : try {
230 1 : args.onResultsReady?.();
231 1 : } catch {
232 : // Scheduler notifications must not interrupt query execution.
233 1 : }
234 1 : };
235 1 : traced.info("Executing query", {
236 1 : sessionId,
237 1 : "query": queryId.toString(),
238 : // Where a statement of a plan of session scripts comes from, e.g. `bar.sql 2/3 · inlines foo.sql`
239 1 : ...(args.metadata.planned != null ? { "planned": args.metadata.title ?? '', "node": args.metadata.planned.node.toString() } : {}),
240 1 : "text": args.query
241 1 : }, LOG_CTX);
242 :
243 : // The cache of the app holds Parquet files that the compute database writes and reads
244 1 : const resultCache = queryResultCache;
245 1 : let cacheHash: string | null = null;
246 1 : if (args.cacheable && args.cacheSignature && resultCache != null && supportsParquet(computeDb)) {
247 1 : cacheHash = await computeQueryCacheKeyForConnection(conn.details, args.cacheSignature);
248 : }
249 1 : const parquetDb = computeDb as ParquetDatabase;
250 :
251 : // Cache read path: on a hit, read the Parquet file and drive the state machine as if
252 : // the result had just streamed in, skipping the backend entirely.
253 1 : const serveFromCache = async (): Promise<arrow.Table | null> => {
254 1 : if (cacheHash == null || args.refreshCache) return null;
255 1 : let cached: QueryResultCacheEntry | null = null;
256 1 : try {
257 1 : cached = await resultCache!.lookup(cacheHash);
258 1 : } catch (e: any) {
259 1 : traced.warn("Failed to read query cache", { query: queryId.toString(), error: stringifyError(e) }, LOG_CTX);
260 1 : }
261 : // A user cancel during the async cache read behaves like any other cancel
262 1 : initialState.cancellation.signal.throwIfAborted();
263 1 : if (cached == null) return null;
264 1 : let table: arrow.Table;
265 1 : try {
266 1 : table = await resultCache!.readTable(parquetDb, cached.name);
267 1 : } catch (e: any) {
268 : // An unreadable entry is a miss
269 1 : traced.warn("Failed to read cached query result", { query: queryId.toString(), error: stringifyError(e) }, LOG_CTX);
270 1 : return null;
271 1 : }
272 1 : initialState.cancellation.signal.throwIfAborted();
273 : // Record the access so eviction sees this as a recently-used (LRU) entry. This
274 : // bumps only the empty `.last_access` marker, never the payload, so it is cheap and
275 : // leaves the payload's "cached at" write time intact. Best-effort: a failure here
276 : // must never surface into the query path (it only means the entry looks colder to
277 : // eviction than it is).
278 1 : void resultCache!.touch(cached.name).catch((e: any) => {
279 1 : traced.warn("Failed to record query cache access", { query: queryId.toString(), error: stringifyError(e) }, LOG_CTX);
280 1 : });
281 1 : traced.info("Served query from cache", {
282 1 : sessionId,
283 1 : "query": queryId.toString(),
284 1 : "numRows": table.numRows.toString(),
285 1 : "numCols": table.numCols.toString(),
286 1 : "cachedAt": new Date(cached.cachedAtMs).toISOString(),
287 1 : }, LOG_CTX);
288 : // No live stream, so synthesize empty metadata and zeroed stream metrics.
289 1 : queryTracker.dispatch({
290 1 : type: QUERY_RECEIVED_ALL_BATCHES,
291 1 : value: [queryId, table, new Map<string, string>(), createQueryResponseStreamMetrics()],
292 1 : });
293 : // Record the cache key and the entry's write time so the UI can show how old the
294 : // cached result is and offer to delete it.
295 1 : queryTracker.dispatch({
296 1 : type: QUERY_CACHE_RECORDED,
297 1 : value: [queryId, cacheHash, true, cached.cachedAtMs],
298 1 : });
299 1 : return table;
300 1 : };
301 : // Compute all table summaries of the result
302 1 : const analyzeResults = async (table: arrow.Table) => {
303 1 : if (!args.analyzeResults) return;
304 1 : queryTracker.dispatch({ type: QUERY_PROCESSING_RESULTS, value: [queryId] });
305 1 : try {
306 1 : await measurePostprocessing(async () => {
307 1 : await yieldToBrowser();
308 1 : await analyzeTable(queryId, table, computeDispatch, computeDb, traced, connectionId);
309 1 : initialState.cancellation.signal.throwIfAborted();
310 1 : });
311 1 : } catch (e) {
312 : // The previous result of the owner stays
313 1 : computeDispatch({ type: DELETE_COMPUTATION, value: [queryId] });
314 1 : throw e;
315 1 : }
316 : // The result replaces the previous result of its owner, or a later result replaced it already
317 1 : const owner = args.metadata.resultOwner;
318 1 : const unused = owner == null ? null : resultComputations.current.retain(owner, queryId);
319 1 : if (unused != null) {
320 1 : computeDispatch({ type: DELETE_COMPUTATION, value: [unused] });
321 : }
322 1 : queryTracker.dispatch({ type: QUERY_PROCESSED_RESULTS, value: [queryId] });
323 1 : };
324 : // Cache write path: after a successful miss, store the result for next time. Fire-and-forget
325 : // (not awaited) so a large eviction scan never stalls the caller's promise, and never fatal —
326 : // a quota/permission failure just logs.
327 1 : const storeInCache = (table: arrow.Table | null) => {
328 1 : if (cacheHash == null || table == null) return;
329 : // Hyper writes the file into the mounted cache, the new file stays pinned until it is complete
330 1 : const write = async () => {
331 1 : const name = resultCache!.newFileName(cacheHash!);
332 1 : const release = await resultCache!.pin(name);
333 1 : try {
334 1 : await resultCache!.writeTable(parquetDb, table, name);
335 1 : } finally {
336 1 : await release();
337 1 : }
338 1 : };
339 1 : void write().then(() => {
340 : // The write landed: record the key (but not servedFromCache — this run hit the
341 : // backend) so the UI can offer to delete the freshly-cached entry. The "cached at"
342 : // time is the write we just made; a later hit reads the precise mtime from disk.
343 1 : queryTracker.dispatch({
344 1 : type: QUERY_CACHE_RECORDED,
345 1 : value: [queryId, cacheHash!, false, null],
346 1 : });
347 1 : }).catch((e: any) => {
348 1 : traced.warn("Failed to write query cache", { query: queryId.toString(), error: stringifyError(e) }, LOG_CTX);
349 1 : });
350 1 : };
351 :
352 : // The catalog follows the DDL statements, a full refresh can be expensive
353 1 : const applyCatalogChanges = (changes: core.buffers.catalog.CatalogChangeT[]) => {
354 1 : connDispatch(connectionId, { type: APPLY_CATALOG_CHANGES, value: changes });
355 1 : try {
356 1 : args.onCatalogChanged?.();
357 1 : } catch {
358 : // Catalog observers must not interrupt query execution.
359 1 : }
360 1 : };
361 :
362 : // Compute the cache key up front for cacheable queries. This is best-effort: if the
363 : // connection has no recoverable params/signature (e.g. before setup completes) we simply skip
364 : // caching and execute normally. Never let a cache concern surface into the query path.
365 : // XXX Add explicit query preparation here later
366 :
367 : // Execute the query and consume the results
368 1 : let resultStream: QueryExecutionResponseStream | null = null;
369 1 : let table: arrow.Table | null = null;
370 1 : let servedFromCache = false;
371 1 : try {
372 1 : table = await serveFromCache();
373 1 : servedFromCache = table != null;
374 1 : if (!servedFromCache && args.cacheOnly) throw new QueryCacheMissError();
375 :
376 1 : if (!servedFromCache) {
377 1 : tracked.sending();
378 :
379 : // Start the query
380 1 : startQueryExecution();
381 1 : resultStream = await executeConnectionQuery(conn.details, args, initialState.cancellation.signal);
382 1 : if (initialState.cancellation.signal.aborted) {
383 1 : await resultStream?.cancel?.();
384 1 : initialState.cancellation.signal.throwIfAborted();
385 1 : }
386 1 : traced.debug("Received query results", {
387 1 : sessionId,
388 1 : "query": queryId.toString()
389 1 : }, LOG_CTX);
390 :
391 1 : if (resultStream != null) {
392 1 : runtime.resultStream = resultStream;
393 1 : resultStream.setProfileDetail?.(runtime.profileViewers > 0);
394 1 : tracked.running(resultStream);
395 :
396 : // Subscribe to query_status and result messages
397 1 : table = await consumeQueryResponseStream({
398 1 : stream: resultStream,
399 1 : abort: initialState.cancellation.signal,
400 1 : publishResults: args.publishResults !== false,
401 1 : onProgress: progress => {
402 1 : const profileProgress = progress.engineProgress == null
403 1 : ? null
404 0 : : applyEngineProgress(queryId, progress.engineProgress);
405 1 : queryTracker.dispatch({
406 1 : type: QUERY_PROGRESS_UPDATED,
407 1 : value: [queryId, profileProgress == null ? progress : { ...progress, profileProgress }],
408 1 : });
409 1 : },
410 1 : onBatch: (batch, stream) => queryTracker.dispatch({
411 1 : type: QUERY_RECEIVED_BATCH,
412 1 : value: [queryId, batch, stream.getMetrics()],
413 1 : }),
414 1 : onResultsReceived: reportQueryExecution,
415 1 : logger: traced,
416 1 : logContext: { sessionId, queryId, target: LOG_CTX },
417 1 : });
418 :
419 1 : traced.info("Executed query", {
420 1 : sessionId,
421 1 : "query": queryId.toString(),
422 1 : "numRows": table?.numRows.toString(),
423 1 : "numCols": table?.numCols.toString(),
424 1 : "batchesReceived": resultStream.getMetrics().totalBatchesReceived.toString(),
425 1 : "dataBytesReceived": resultStream.getMetrics().totalDataBytesReceived.toString(),
426 1 : }, LOG_CTX);
427 :
428 : // Is there any metadata?
429 1 : const metadata = resultStream.getMetadata();
430 1 : if (table != null) queryTracker.dispatch({
431 1 : type: QUERY_RECEIVED_ALL_BATCHES,
432 1 : value: [queryId, table!, metadata, resultStream!.getMetrics()],
433 1 : });
434 1 : } else {
435 1 : reportQueryExecution();
436 1 : traced.warn("Query returned no results", { sessionId, "query": queryId.toString() }, LOG_CTX);
437 1 : }
438 1 : }
439 1 : } catch (e: any) {
440 1 : if (initialState.cancellation.signal.aborted || e?.name === 'AbortError' || e?.message === 'AbortError') {
441 1 : const cancellationError = e instanceof LoggableException
442 1 : ? e
443 1 : : new LoggableException('Query was cancelled', {}, LOG_CTX);
444 1 : traced.warn("Cancelled query", {
445 1 : query: queryId.toString(),
446 1 : sessionId
447 1 : }, LOG_CTX);
448 1 : tracked.cancel(cancellationError, resultStream?.getMetrics() ?? null);
449 1 : } else {
450 1 : logQueryFailure(traced, e, queryId, sessionId);
451 1 : tracked.fail(e, resultStream?.getMetrics() ?? null);
452 1 : }
453 1 : if (args.throwOnError) {
454 1 : throw e;
455 : }
456 1 : return null;
457 1 : }
458 1 : if (!servedFromCache && args.catalogChanges != null && args.catalogChanges.length > 0) {
459 1 : applyCatalogChanges(args.catalogChanges);
460 : }
461 :
462 1 : reportResultsReady();
463 :
464 1 : if (table != null) {
465 1 : try {
466 1 : await analyzeResults(table);
467 1 : } catch (e: any) {
468 1 : if (initialState.cancellation.signal.aborted || e?.name === 'AbortError') {
469 1 : const cancellationError = e instanceof LoggableException
470 1 : ? e
471 1 : : new LoggableException('Query was cancelled', {}, LOG_CTX);
472 1 : traced.warn("Cancelled query during result processing", {
473 1 : query: queryId.toString(),
474 1 : sessionId,
475 1 : }, LOG_CTX);
476 1 : tracked.cancel(cancellationError, resultStream?.getMetrics() ?? null);
477 1 : if (args.throwOnError) throw e;
478 1 : return null;
479 1 : }
480 1 : const processingError = e instanceof LoggableException
481 1 : ? e
482 1 : : new LoggableException("Query result processing failed", {
483 1 : error: stringifyError(e),
484 1 : }, LOG_CTX);
485 1 : traced.warn("Query result processing failed", {
486 1 : query: queryId.toString(),
487 1 : sessionId,
488 1 : error: processingError.message,
489 1 : ...processingError.keyValues,
490 1 : }, LOG_CTX);
491 1 : tracked.fail(processingError);
492 1 : throw processingError;
493 1 : }
494 : }
495 :
496 1 : if (!servedFromCache) storeInCache(table);
497 1 : return table;
498 1 : },
499 1 : });
500 1 : return execution;
501 :
502 1 : }, [computeDb, connStore, connDispatch, computeDispatch, logger, queryResultCache, applyEngineProgress]);
503 :
504 : // Allocate the next query id and start the execution
505 1 : const execute = React.useCallback<QueryExecutor>((connectionId: string, args: QueryExecutionArgs): [number, Promise<arrow.Table | null>] => {
506 1 : const queryId = allocateQueryId();
507 1 : const trace = createTrace();
508 1 : const runtime: QueryExecutionRuntime = {
509 1 : cancellation: new AbortController(),
510 1 : resultStream: null,
511 1 : backendCancellation: null,
512 1 : execution: null,
513 1 : profileViewers: 0,
514 1 : };
515 : // Every abort stops the query at the backend, a stream that arrives later is cancelled when it arrives
516 1 : runtime.cancellation.signal.addEventListener('abort', () => {
517 1 : const stream = runtime.resultStream;
518 1 : if (stream?.cancel == null) return;
519 1 : runtime.backendCancellation = stream.cancel().catch((e: any) => {
520 1 : logger.warn('Failed to cancel query at the backend', {
521 1 : query: queryId.toString(),
522 1 : connectionId,
523 1 : error: stringifyError(e),
524 1 : }, LOG_CTX);
525 1 : });
526 1 : }, { once: true });
527 1 : const cancelFromCaller = () => runtime.cancellation.abort(args.abortSignal?.reason);
528 1 : if (args.abortSignal?.aborted) {
529 1 : cancelFromCaller();
530 : } else {
531 1 : args.abortSignal?.addEventListener('abort', cancelFromCaller, { once: true });
532 : }
533 1 : queryRuntimes.current.set(queryId, runtime);
534 1 : const logObserver = args.onLog == null ? null : (record: LogRecord) => {
535 1 : try {
536 1 : args.onLog?.(record.message);
537 1 : } catch {
538 : // A progress surface must never interrupt query execution or logger delivery.
539 1 : }
540 1 : };
541 1 : if (logObserver != null) logger.buffer.subscribeTrace(trace.traceId, logObserver);
542 1 : const execution = executeImpl(connectionId, args, queryId, runtime, trace);
543 1 : runtime.execution = execution;
544 1 : const removeRuntime = () => {
545 1 : args.abortSignal?.removeEventListener('abort', cancelFromCaller);
546 1 : queryRuntimes.current.delete(queryId);
547 1 : progressModels.current.release(queryId);
548 1 : if (logObserver != null) logger.buffer.unsubscribeTrace(trace.traceId, logObserver);
549 1 : };
550 1 : void execution.then(removeRuntime, removeRuntime);
551 1 : return [queryId, execution];
552 1 : }, [executeImpl, logger]);
553 :
554 1 : const cancel = React.useCallback<CancelQuery>(async (connectionId, queryId) => {
555 1 : const runtime = queryRuntimes.current.get(queryId);
556 1 : if (runtime == null) return;
557 1 : runtime.cancellation.abort();
558 1 : await runtime.backendCancellation;
559 1 : await runtime.execution?.catch(() => null);
560 1 : }, []);
561 :
562 1 : const watchProfile = React.useCallback<WatchQueryProfile>(queryId => {
563 1 : const update = (delta: number) => {
564 1 : const runtime = queryRuntimes.current.get(queryId);
565 1 : if (runtime == null) return;
566 1 : runtime.profileViewers = Math.max(0, runtime.profileViewers + delta);
567 1 : runtime.resultStream?.setProfileDetail?.(runtime.profileViewers > 0);
568 1 : };
569 1 : update(1);
570 1 : return () => update(-1);
571 1 : }, []);
572 :
573 1 : const value = React.useMemo(() => ({ execute, cancel, watchProfile }), [execute, cancel, watchProfile]);
574 :
575 1 : return (
576 1 : <EXECUTOR_CTX.Provider value={value}>
577 1 : {props.children}
578 1 : </EXECUTOR_CTX.Provider>
579 : );
580 1 : }
|
