Line data    Source code

       1              : import * as React from 'react';
       2              : 
       3              : import { ComputeQueryExecutionProvider } from '../compute/compute_query_execution.js';
       4              : import type { ComputeQueryExecution } from '../compute/computation_tasks.js';
       5              : import { useComputationStore } from '../compute/computation_registry.js';
       6              : import { computationTaskTitle, type TaskVariant } from '../compute/computation_types.js';
       7              : import { QueryType } from '../query/query_execution_state.js';
       8              : import { executeTrackedQuery } from '../query/tracked_query_execution.js';
       9              : import { useConnectionDispatch } from '../connections/connection_registry.js';
      10              : import { sqlframeClientContext } from '../connections/hyper/hyper_log_metrics.js';
      11              : 
      12            1 : const LOG_CTX = 'scheduler';
      13              : 
      14              : /// Track the queries of computations in the query history of the connection that ran the query of the table
      15            1 : export function SessionComputeQueryExecutionProvider(props: React.PropsWithChildren) {
      16              :     // Resolve the connection when the task starts, subscribing to the computations would re-create the
      17              :     // execution factory and re-render the scheduler for every task step.
      18            1 :     const computations = useComputationStore();
      19            1 :     const dispatch = useConnectionDispatch();
      20            1 :     const createExecution = React.useCallback((task: TaskVariant) => {
      21            1 :         const computation = computations.current.tableComputations[task.value.tableId];
      22            1 :         const connectionId = computation?.connectionId ?? null;
      23            1 :         const execution: ComputeQueryExecution = (query, execute) => {
      24            1 :             if (connectionId == null) return execute(computation?.clientContext ?? null);
      25            1 :             return executeTrackedQuery({
      26            1 :                 query,
      27            1 :                 tracker: { dispatch: action => dispatch(connectionId, action) },
      28            1 :                 metadata: {
      29            1 :                     queryType: QueryType.INTERNAL_SQLFRAME,
      30            1 :                     title: computationTaskTitle(task),
      31            1 :                     description: null,
      32            1 :                     issuer: 'SQLFrame',
      33            1 :                     userProvided: false,
      34            1 :                 },
      35              :                 // Attribute the query to the connection in the Hyper log, it runs on a pooled compute connection
      36            1 :                 execute: () => execute(sqlframeClientContext(connectionId)),
      37            1 :                 errorTarget: LOG_CTX,
      38            1 :             });
      39            1 :         };
      40            1 :         return execution;
      41            1 :     }, [computations, dispatch]);
      42              : 
      43            1 :     return (
      44            1 :         <ComputeQueryExecutionProvider createExecution={createExecution}>
      45            1 :             {props.children}
      46            1 :         </ComputeQueryExecutionProvider>
      47              :     );
      48            1 : }