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