Line data    Source code

       1              : import {
       2              :     EXECUTE_QUERY,
       3              :     QUERY_CANCELLED,
       4              :     QUERY_FAILED,
       5              :     QUERY_RUNNING,
       6              :     QUERY_SENDING,
       7              :     QUERY_SUCCEEDED,
       8              :     type QueryExecutionAction,
       9              :     createQueryExecutionState,
      10              :     type QueryExecutionMetrics,
      11              :     type QueryExecutionResponseStream,
      12              :     type QueryExecutionState,
      13              :     type QueryExecutionTracker,
      14              :     type QueryMetadata,
      15              : } from './query_execution_state.js';
      16              : import { LoggableException, stringifyError } from '../platform/logger/logger.js';
      17              : import { createTrace, type TraceContext } from '../platform/logger/trace_context.js';
      18              : 
      19            1 : let NEXT_QUERY_ID = 1;
      20              : 
      21            1 : export function allocateQueryId(): number {
      22            1 :     return NEXT_QUERY_ID++;
      23            1 : }
      24              : 
      25              : export interface TrackedQueryExecutionContext {
      26              :     readonly queryId: number;
      27              :     readonly trace: TraceContext;
      28              :     readonly cancellation: AbortController;
      29              :     readonly initialState: QueryExecutionState;
      30              :     readonly isDone: () => boolean;
      31              :     readonly dispatch: (action: QueryExecutionAction) => void;
      32              :     readonly sending: () => void;
      33              :     readonly running: (stream?: QueryExecutionResponseStream | null) => void;
      34              :     readonly succeed: () => void;
      35              :     readonly fail: (error: unknown, metrics?: QueryExecutionMetrics | null) => void;
      36              :     readonly cancel: (error: unknown, metrics?: QueryExecutionMetrics | null) => void;
      37              : }
      38              : 
      39              : export interface TrackedQueryExecutionArgs<T> {
      40              :     query: string;
      41              :     metadata: QueryMetadata;
      42              :     tracker: QueryExecutionTracker;
      43              :     execute: (context: TrackedQueryExecutionContext) => Promise<T>;
      44              :     queryId?: number;
      45              :     trace?: TraceContext;
      46              :     cancellation?: AbortController;
      47              :     autoStart?: boolean;
      48              :     errorTarget?: string;
      49              : }
      50              : 
      51            1 : function asLoggableException(error: unknown, target: string): LoggableException {
      52            1 :     return error instanceof LoggableException
      53            1 :         ? error
      54            1 :         : new LoggableException(stringifyError(error), {}, target);
      55            1 : }
      56              : 
      57            1 : export async function executeTrackedQuery<T>(args: TrackedQueryExecutionArgs<T>): Promise<T> {
      58            1 :     const queryId = args.queryId ?? allocateQueryId();
      59            1 :     const trace = args.trace ?? createTrace();
      60            1 :     const cancellation = args.cancellation ?? new AbortController();
      61            1 :     const initialState = createQueryExecutionState(
      62            1 :         queryId,
      63            1 :         trace.traceId,
      64            1 :         args.query,
      65            1 :         args.metadata,
      66            1 :         cancellation,
      67            1 :     );
      68            1 :     let done = false;
      69            1 :     const dispatch = (action: QueryExecutionAction) => args.tracker.dispatch(action);
      70            1 :     const finish = (action: QueryExecutionAction) => {
      71            1 :         if (done) return;
      72            1 :         done = true;
      73            1 :         dispatch(action);
      74            1 :     };
      75            1 :     const errorTarget = args.errorTarget ?? 'query_execution';
      76            1 :     const context: TrackedQueryExecutionContext = {
      77            1 :         queryId,
      78            1 :         trace,
      79            1 :         cancellation,
      80            1 :         initialState,
      81            1 :         isDone: () => done,
      82            1 :         dispatch,
      83            1 :         sending: () => dispatch({ type: QUERY_SENDING, value: [queryId] }),
      84            1 :         running: stream => dispatch({ type: QUERY_RUNNING, value: [queryId, stream ?? null] }),
      85            1 :         succeed: () => finish({ type: QUERY_SUCCEEDED, value: [queryId] }),
      86            1 :         fail: (error, metrics = null) => finish({ type: QUERY_FAILED, value: [queryId, asLoggableException(error, errorTarget), metrics] }),
      87            1 :         cancel: (error, metrics = null) => finish({ type: QUERY_CANCELLED, value: [queryId, asLoggableException(error, errorTarget), metrics] }),
      88            1 :     };
      89              : 
      90            1 :     dispatch({ type: EXECUTE_QUERY, value: [queryId, initialState] });
      91            1 :     if (args.autoStart !== false) {
      92            1 :         context.sending();
      93            1 :         context.running();
      94            1 :     }
      95              : 
      96            1 :     try {
      97            1 :         cancellation.signal.throwIfAborted();
      98            1 :         const result = await args.execute(context);
      99            1 :         context.succeed();
     100            1 :         return result;
     101            1 :     } catch (error: any) {
     102            1 :         if (!done) {
     103            1 :             if (cancellation.signal.aborted || error?.name === 'AbortError' || error?.message === 'AbortError') {
     104            1 :                 context.cancel(error);
     105              :             } else {
     106            1 :                 context.fail(error);
     107              :             }
     108            1 :         }
     109            1 :         throw error;
     110            1 :     }
     111            1 : }