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 { reportShellTiming, SHELL_QUERY_METADATA } from '../shell/shell_result.js';
      14              : 
      15              : interface EmbeddedDatabaseShellEnvironmentOptions {
      16              :     prepareResult?: (queryId: number, table: arrow.Table) => void | Promise<void>;
      17              :     /// Called after a statement ran or failed, before its result shows and the prompt returns
      18              :     afterStatement?: (query: string) => 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 :                         let table: arrow.Table;
      45            1 :                         try {
      46            1 :                             table = await connection.query(query, cancellation.signal);
      47            1 :                         } catch (error) {
      48            1 :                             await options.afterStatement?.(query);
      49            1 :                             throw error;
      50            1 :                         }
      51            1 :                         reportTiming('Results', performance.now() - queryStartedAt);
      52            1 :                         cancellation.signal.throwIfAborted();
      53            1 :                         await options.afterStatement?.(query);
      54            1 :                         const postprocessingStartedAt = performance.now();
      55            1 :                         try {
      56            1 :                             const metrics = createQueryResponseStreamMetrics();
      57            1 :                             metrics.totalDataBytesReceived = table.batches.reduce((bytes, batch) => bytes + batch.data.byteLength, 0);
      58            1 :                             metrics.totalBatchesReceived = table.batches.length;
      59            1 :                             metrics.totalRowsReceived = table.numRows;
      60            1 :                             tracked.dispatch({
      61            1 :                                 type: QUERY_RECEIVED_ALL_BATCHES,
      62            1 :                                 value: [tracked.queryId, table, new Map(), metrics],
      63            1 :                             });
      64            1 :                             if (table.numCols === 0) return;
      65            1 :                             if (options.prepareResult != null) {
      66            1 :                                 tracked.dispatch({
      67            1 :                                     type: QUERY_PROCESSING_RESULTS,
      68            1 :                                     value: [tracked.queryId],
      69            1 :                                 });
      70            1 :                                 await options.prepareResult(tracked.queryId, table);
      71            1 :                                 cancellation.signal.throwIfAborted();
      72            1 :                                 tracked.dispatch({
      73            1 :                                     type: QUERY_PROCESSED_RESULTS,
      74            1 :                                     value: [tracked.queryId],
      75            1 :                                 });
      76            1 :                             }
      77            1 :                             onResult?.(tracked.queryId, table.numRows);
      78            1 :                         } finally {
      79            1 :                             reportTiming('Postprocess', performance.now() - postprocessingStartedAt);
      80            1 :                         }
      81            1 :                     },
      82            1 :                 });
      83            1 :             } finally {
      84            1 :                 signal?.removeEventListener('abort', abort);
      85            1 :             }
      86            1 :         },
      87            1 :     };
      88            1 : }