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