Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import type { DashQLShellCommand } from './api.js';
       4              : import { QueryType, type QueryMetadata } from '../query/query_execution_state.js';
       5              : import { detectPlanEngine } from '../views/query_profile/plan_engine.js';
       6              : 
       7              : export const SHELL_AUTO_OVERLAY_ROW_LIMIT = 100;
       8            1 : const TABLE_FRAME_WIDTH_PER_COLUMN = 3;
       9              : 
      10              : export type ShellOutputMode = 'auto' | 'ui' | 'term' | 'off';
      11              : 
      12              : /// The result of a query whose output the shell doesn't print
      13            1 : export const EMPTY_RESULT_IPC = arrow.tableToIPC(arrow.tableFromArrays({}), 'file');
      14              : 
      15              : /// The metadata of the queries that users enter in the shell
      16            1 : export const SHELL_QUERY_METADATA: QueryMetadata = {
      17            1 :     queryType: QueryType.USER_PROVIDED,
      18            1 :     title: 'Shell Query',
      19            1 :     description: null,
      20            1 :     issuer: 'DashQL Shell',
      21            1 :     userProvided: true,
      22            1 : };
      23              : 
      24              : /// Report a timing to an optional observer, observers must not interrupt the query execution
      25            1 : export function reportShellTiming(onTiming: ((label: string, elapsedMs: number) => void) | undefined, label: string, elapsedMs: number): void {
      26            1 :     try {
      27            1 :         onTiming?.(label, elapsedMs);
      28            1 :     } catch {
      29              :         // Timing observers must not interrupt query execution.
      30            0 :     }
      31            1 : }
      32              : 
      33            1 : export function getPlanResultText(table: arrow.Table | null): string | null {
      34            1 :     if (table == null || table.numRows !== 1 || table.numCols !== 1) return null;
      35            1 :     const value = table.getChildAt(0)?.get(0);
      36            1 :     if (typeof value !== 'string') return null;
      37            1 :     if (detectPlanEngine(value) == null) return null;
      38            1 :     try {
      39            1 :         const parsed = JSON.parse(value);
      40            1 :         if (parsed == null || typeof parsed !== 'object' || Array.isArray(parsed)) return null;
      41              :         // A Hyper plan has a root operator, a Trino plan maps fragment ids to plan nodes
      42            1 :         const isHyperPlan = typeof parsed.operator === 'string';
      43            1 :         const isTrinoPlan = Object.values(parsed).every(node =>
      44            1 :             node != null && typeof node === 'object' && typeof (node as { name?: unknown }).name === 'string');
      45            1 :         return isHyperPlan || isTrinoPlan ? value : null;
      46            1 :     } catch {
      47            0 :         return null;
      48            0 :     }
      49            1 : }
      50              : 
      51            1 : function textWidth(value: unknown): number {
      52            1 :     if (value == null) return 0;
      53            1 :     return String(value).split(/\r?\n/, -1).reduce((width, line) => Math.max(width, Array.from(line).length), 0);
      54            1 : }
      55              : 
      56            1 : export function estimateTerminalTableWidth(table: arrow.Table): number {
      57            1 :     if (table.numCols === 0) return 0;
      58            1 :     let width = TABLE_FRAME_WIDTH_PER_COLUMN * table.numCols + 1;
      59            1 :     for (let columnIndex = 0; columnIndex < table.numCols; ++columnIndex) {
      60            1 :         const column = table.getChildAt(columnIndex);
      61            1 :         let columnWidth = textWidth(table.schema.fields[columnIndex]?.name ?? '');
      62            1 :         if (column != null) {
      63            1 :             for (let rowIndex = 0; rowIndex < table.numRows; ++rowIndex) {
      64            1 :                 columnWidth = Math.max(columnWidth, textWidth(column.get(rowIndex)));
      65              :             }
      66              :         }
      67            1 :         width += Math.max(columnWidth, 1);
      68            1 :     }
      69            1 :     return width;
      70            1 : }
      71              : 
      72            1 : export function shouldShowResultUI(mode: ShellOutputMode, table: arrow.Table, terminalColumns: number): boolean {
      73            1 :     return mode === 'ui'
      74            1 :         || (mode === 'auto' && (
      75            1 :             getPlanResultText(table) != null
      76              :             ||
      77            1 :             table.numRows > SHELL_AUTO_OVERLAY_ROW_LIMIT
      78            1 :             || estimateTerminalTableWidth(table) > Math.max(terminalColumns, 1)
      79              :         ));
      80            1 : }
      81              : 
      82            1 : export function createShellOutputCommand(
      83            1 :     getMode: () => ShellOutputMode,
      84            1 :     setMode: (mode: ShellOutputMode) => void,
      85            1 : ): DashQLShellCommand {
      86            1 :     return [
      87            1 :         'output',
      88            1 :         'Set query output: auto, ui, term, or off',
      89            1 :         args => {
      90            1 :             if (args.length > 1 || (args.length === 1 && !['auto', 'ui', 'term', 'off'].includes(args[0]))) {
      91            1 :                 throw new Error('usage: .output [auto|ui|term|off]');
      92              :             }
      93            1 :             if (args.length === 1) setMode(args[0] as ShellOutputMode);
      94            1 :             const mode = getMode();
      95            1 :             return mode === 'auto'
      96            1 :                 ? `Query output: auto (UI when results exceed ${SHELL_AUTO_OVERLAY_ROW_LIMIT} rows or terminal width)`
      97            1 :                 : `Query output: ${mode}`;
      98            1 :         },
      99            1 :     ];
     100            1 : }