Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import { CancelQuery, QueryExecutor } from '../../connections/query_executor.js';
       4              : import type { DashQLShellEnvironment } from '../../shell/api.js';
       5              : import {
       6              :     EMPTY_RESULT_IPC,
       7              :     reportShellTiming,
       8              :     SHELL_QUERY_METADATA,
       9              :     shouldShowResultUI,
      10              :     type ShellOutputMode,
      11              : } from '../../shell/shell_result.js';
      12              : 
      13            1 : export function createSessionShellEnvironment(
      14            1 :     sessionId: string,
      15            1 :     connectionId: string,
      16            1 :     executeQuery: QueryExecutor,
      17            1 :     cancelQuery: CancelQuery,
      18            1 :     getOutputMode: () => ShellOutputMode = () => 'auto',
      19            1 :     getTerminalColumns: () => number = () => 100,
      20            1 :     releaseResult: (queryId: number) => void = () => { },
      21            1 : ): DashQLShellEnvironment {
      22            1 :     return {
      23            1 :         async executeQuery(query, signal, onProgress, onResult, onTiming) {
      24            1 :             const reportTiming = (label: string, elapsedMs: number) => reportShellTiming(onTiming, label, elapsedMs);
      25            1 :             const [queryId, execution] = executeQuery(connectionId, {
      26            1 :                 query,
      27            1 :                 analyzeResults: true,
      28            1 :                 cacheable: false,
      29            1 :                 throwOnError: true,
      30            1 :                 onLog: onProgress,
      31            1 :                 onTiming: reportTiming,
      32            1 :                 metadata: { ...SHELL_QUERY_METADATA, sessionId },
      33            1 :             });
      34            1 :             const abort = () => cancelQuery(connectionId, queryId);
      35            1 :             if (signal?.aborted) {
      36            0 :                 abort();
      37              :             } else {
      38            1 :                 signal?.addEventListener('abort', abort, { once: true });
      39              :             }
      40            1 :             try {
      41            1 :                 const table = await execution;
      42            1 :                 if (table == null) {
      43            0 :                     if (signal?.aborted) throw new DOMException('Query was cancelled', 'AbortError');
      44            0 :                     throw new Error('Query failed without an error');
      45            0 :                 }
      46            1 :                 const outputStartedAt = performance.now();
      47            1 :                 try {
      48            1 :                     const outputMode = getOutputMode();
      49            1 :                     if (table.numCols > 0 && shouldShowResultUI(outputMode, table, getTerminalColumns())) {
      50            1 :                         onResult?.(queryId, table.numRows);
      51            1 :                         return EMPTY_RESULT_IPC;
      52            1 :                     }
      53              :                     // The terminal shows the result as text, its computation is not needed
      54            1 :                     releaseResult(queryId);
      55            1 :                     if (outputMode === 'off') return EMPTY_RESULT_IPC;
      56            1 :                     return arrow.tableToIPC(table, 'file');
      57            1 :                 } finally {
      58            1 :                     reportTiming('Postprocess', performance.now() - outputStartedAt);
      59            1 :                 }
      60            1 :             } finally {
      61            1 :                 signal?.removeEventListener('abort', abort);
      62            1 :             }
      63            1 :         },
      64            1 :     };
      65            1 : }