Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import { AsyncConsumerLambdas } from '../utils/async_consumer.js';
       4              : import type { TracedLogger } from '../platform/logger/logger.js';
       5              : import type { QueryExecutionProgress, QueryExecutionResponseStream } from '../query/query_execution_state.js';
       6              : import { executeConnectionClientFunction, type ClientFunctionServices } from './connection_client_functions.js';
       7              : import { connectorOfDetails } from './connector_descriptor.js';
       8              : import type { ConnectionStateDetailsVariant } from './connection_state_details.js';
       9              : import type { QueryExecutionArgs } from './query_execution_args.js';
      10              : 
      11              : /// Execute a query on the channel of a connection, a client function call as requests to its service
      12            1 : export function executeConnectionQuery(
      13            1 :     details: ConnectionStateDetailsVariant,
      14            1 :     args: QueryExecutionArgs,
      15            1 :     services: ClientFunctionServices,
      16            1 :     abort?: AbortSignal,
      17            1 : ): Promise<QueryExecutionResponseStream> {
      18            1 :     if (args.clientFunctionCall != null) {
      19            1 :         return executeConnectionClientFunction(details, args.clientFunctionCall, services, abort);
      20              :     }
      21            1 :     return connectorOfDetails(details).executeQuery(details.value, args, abort);
      22            1 : }
      23              : 
      24              : interface ConsumeQueryResponseStreamArgs {
      25              :     stream: QueryExecutionResponseStream;
      26              :     abort?: AbortSignal;
      27              :     publishResults: boolean;
      28              :     onProgress: (progress: QueryExecutionProgress) => void;
      29              :     onBatch?: (batch: arrow.RecordBatch, stream: QueryExecutionResponseStream) => void;
      30              :     onResultsReceived?: () => void;
      31              :     logger?: TracedLogger;
      32              :     logContext?: {
      33              :         sessionId: string;
      34              :         queryId: number;
      35              :         target: string;
      36              :     };
      37              : }
      38              : 
      39            1 : export async function consumeQueryResponseStream(args: ConsumeQueryResponseStreamArgs): Promise<arrow.Table | null> {
      40            1 :     const batches: arrow.RecordBatch[] = [];
      41            1 :     const consumeProgress = new AsyncConsumerLambdas<QueryExecutionResponseStream, QueryExecutionProgress>(
      42            1 :         (_stream, progress) => args.onProgress(progress),
      43            1 :     );
      44            1 :     const consumeBatches = new AsyncConsumerLambdas<QueryExecutionResponseStream, arrow.RecordBatch>(
      45            1 :         (stream, batch) => {
      46            1 :             if (!args.publishResults) return;
      47            1 :             batches.push(batch);
      48            1 :             if (args.logger && args.logContext) {
      49            1 :                 args.logger.debug("Received result batch", {
      50            1 :                     sessionId: args.logContext.sessionId,
      51            1 :                     query: args.logContext.queryId.toString(),
      52            1 :                     batchColumns: batch.numCols.toString(),
      53            1 :                     batchRows: batch.numRows.toString(),
      54            1 :                 }, args.logContext.target);
      55              :             }
      56            1 :             args.onBatch?.(batch, stream);
      57            1 :         },
      58            1 :     );
      59              : 
      60            1 :     await args.stream.produce(consumeBatches, consumeProgress, args.abort);
      61            1 :     args.onResultsReceived?.();
      62            1 :     if (!args.publishResults) return null;
      63              : 
      64            1 :     const schema = batches.length > 0
      65            1 :         ? batches[0].schema
      66            1 :         : await args.stream.getSchema() ?? new arrow.Schema();
      67            1 :     return new arrow.Table(schema, batches);
      68            1 : }