Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import {
       4              :     QUERY_PROCESSED_RESULTS,
       5              :     QUERY_PROCESSING_RESULTS,
       6              :     QUERY_RECEIVED_ALL_BATCHES,
       7              :     createQueryResponseStreamMetrics,
       8              :     type QueryExecutionTracker,
       9              : } from '../query/query_execution_state.js';
      10              : import type { EmbeddedConnection } from '../platform/database/embedded_database.js';
      11              : import { DashQLShellEnvironment } from '../shell/api.js';
      12              : import { executeTrackedQuery } from '../query/tracked_query_execution.js';
      13              : import { EMPTY_RESULT_IPC, reportShellTiming, SHELL_QUERY_METADATA, shouldShowResultUI, type ShellOutputMode } from '../shell/shell_result.js';
      14              : 
      15              : interface EmbeddedDatabaseShellEnvironmentOptions {
      16              :     getOutputMode?: () => ShellOutputMode;
      17              :     getTerminalColumns?: () => number;
      18              :     prepareResult?: (queryId: number, table: arrow.Table) => void | Promise<void>;
      19              : }
      20              : 
      21            1 : export function createEmbeddedDatabaseShellEnvironment(
      22            1 :     connection: EmbeddedConnection,
      23            1 :     queryExecutions?: QueryExecutionTracker,
      24            1 :     options: EmbeddedDatabaseShellEnvironmentOptions = {},
      25            1 : ): DashQLShellEnvironment {
      26            1 :     return {
      27            1 :         executeQuery: async (query, signal, onProgress, onResult, onTiming) => {
      28            1 :             const reportTiming = (label: string, elapsedMs: number) => reportShellTiming(onTiming, label, elapsedMs);
      29            1 :             const cancellation = new AbortController();
      30            1 :             const abort = () => cancellation.abort();
      31            1 :             if (signal?.aborted) abort();
      32            1 :             else signal?.addEventListener('abort', abort, { once: true });
      33            1 :             const tracker = queryExecutions ?? { dispatch: () => {} };
      34            1 :             try {
      35            1 :                 return await executeTrackedQuery({
      36            1 :                     query,
      37            1 :                     tracker,
      38            1 :                     cancellation,
      39            1 :                     errorTarget: 'standalone_shell',
      40            1 :                     metadata: { ...SHELL_QUERY_METADATA },
      41            1 :                     execute: async tracked => {
      42            1 :                         onProgress?.('Executing query');
      43            1 :                         const queryStartedAt = performance.now();
      44            1 :                         const table = await connection.query(query, cancellation.signal);
      45            1 :                         reportTiming('Results', performance.now() - queryStartedAt);
      46            1 :                         cancellation.signal.throwIfAborted();
      47            1 :                         const postprocessingStartedAt = performance.now();
      48            1 :                         try {
      49            1 :                             const metrics = createQueryResponseStreamMetrics();
      50            1 :                             metrics.totalDataBytesReceived = table.batches.reduce((bytes, batch) => bytes + batch.data.byteLength, 0);
      51            1 :                             metrics.totalBatchesReceived = table.batches.length;
      52            1 :                             metrics.totalRowsReceived = table.numRows;
      53            1 :                             tracked.dispatch({
      54            1 :                                 type: QUERY_RECEIVED_ALL_BATCHES,
      55            1 :                                 value: [tracked.queryId, table, new Map(), metrics],
      56            1 :                             });
      57            1 :                             const outputMode = options.getOutputMode?.() ?? 'auto';
      58            1 :                             const showUI = shouldShowResultUI(
      59            1 :                                 outputMode,
      60            1 :                                 table,
      61            1 :                                 options.getTerminalColumns?.() ?? 100,
      62              :                             );
      63            1 :                             if (showUI && table.numCols > 0 && onResult != null) {
      64            1 :                                 if (options.prepareResult != null) {
      65            1 :                                     tracked.dispatch({
      66            1 :                                         type: QUERY_PROCESSING_RESULTS,
      67            1 :                                         value: [tracked.queryId],
      68            1 :                                     });
      69            1 :                                     await options.prepareResult(tracked.queryId, table);
      70            1 :                                     cancellation.signal.throwIfAborted();
      71            1 :                                     tracked.dispatch({
      72            1 :                                         type: QUERY_PROCESSED_RESULTS,
      73            1 :                                         value: [tracked.queryId],
      74            1 :                                     });
      75            1 :                                 }
      76            1 :                                 onResult(tracked.queryId, table.numRows);
      77            1 :                                 return EMPTY_RESULT_IPC;
      78            1 :                             }
      79            1 :                             if (outputMode === 'off' || table.numCols === 0) return EMPTY_RESULT_IPC;
      80            1 :                             return arrow.tableToIPC(table, 'stream');
      81            1 :                         } finally {
      82            1 :                             reportTiming('Postprocess', performance.now() - postprocessingStartedAt);
      83            1 :                         }
      84            1 :                     },
      85            1 :                 });
      86            1 :             } finally {
      87            1 :                 signal?.removeEventListener('abort', abort);
      88            1 :             }
      89            1 :         },
      90            1 :     };
      91            1 : }