Adding handling for shift + arrow keys to select data and fixing result grid summary (#20200)

Author: aasimkhan30

localization/l10n/bundle.l10n.json

@@ -1486,6 +1486,12 @@
       "{0} is the average, {1} is the count, {2} is the distinct count, {3} is the max, {4} is the min, {5} is the null count, {6} is the sum"
     ]
   },
+  "You have selected data across {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?/{0} is the number of rows to fetch summary statistics for": {
+    "message": "You have selected data across {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?",
+    "comment": [
+      "{0} is the number of rows to fetch summary statistics for"
+    ]
+  },
   "An error occurred while retrieving rows: {0}/{0} is the error message": {
     "message": "An error occurred while retrieving rows: {0}",
     "comment": [

localization/xliff/vscode-mssql.xlf

@@ -3221,6 +3221,10 @@
     <trans-unit id="++CODE++867d2664362bce58e979ee3de5b436d12ab50da730929367290fd6b9043b2a90">
       <source xml:lang="en">You are not connected to any database.</source>
     </trans-unit>
+    <trans-unit id="++CODE++9b358aeeb5265d934737d601151f0982e0b833eaa6abb54c40be62ed494882b8">
+      <source xml:lang="en">You have selected data across {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?</source>
+      <note>{0} is the number of rows to fetch summary statistics for</note>
+    </trans-unit>
     <trans-unit id="++CODE++437f6eb2a46a074059cbbd924d82e7b72146b3bbb13c028906a45bbffdf69e3f">
       <source xml:lang="en">You must be signed into Azure in order to browse SQL databases.</source>
     </trans-unit>

src/constants/locConstants.ts

@@ -949,6 +949,13 @@ export class QueryResult {
                 "{0} is the average, {1} is the count, {2} is the distinct count, {3} is the max, {4} is the min, {5} is the null count, {6} is the sum",
             ],
         });
+    public static summaryFetchConfirmation = (numRows: number) =>
+        l10n.t({
+            message:
+                "You have selected data across {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?",
+            args: [numRows],
+            comment: ["{0} is the number of rows to fetch summary statistics for"],
+        });
     public static getRowsError = (error: string) =>
         l10n.t({
             message: "An error occurred while retrieving rows: {0}",

src/controllers/queryRunner.ts

@@ -51,6 +51,8 @@ import * as os from "os";
 import { Deferred } from "../protocol";
 import { sendActionEvent } from "../telemetry/telemetry";
 import { TelemetryActions, TelemetryViews } from "../sharedInterfaces/telemetry";
+import { SelectionSummaryStats } from "../sharedInterfaces/queryResult";
+import { calculateSelectionSummaryFromData } from "../queryResult/utils";
 
 export interface IResultSet {
     columns: string[];
@@ -1053,6 +1055,55 @@ export default class QueryRunner {
         await this.writeStringToClipboard(insertIntoString);
     }
 
+    public async generateSelectionSummaryData(
+        selection: ISlickRange[],
+        batchId: number,
+        resultId: number,
+    ): Promise<SelectionSummaryStats> {
+        // Keep copy order deterministic
+        selection.sort((a, b) => a.fromRow - b.fromRow);
+
+        let totalRows = 0;
+        for (let range of selection) {
+            totalRows += range.toRow - range.fromRow + 1;
+        }
+
+        const summaryFetchThreshold =
+            vscode.workspace
+                .getConfiguration()
+                .get<number>(Constants.configInMemoryDataProcessingThreshold) ?? 5000;
+
+        if (totalRows > summaryFetchThreshold) {
+            let confirm = await vscode.window.showInformationMessage(
+                LocalizedConstants.QueryResult.summaryFetchConfirmation(totalRows),
+                { modal: false },
+                LocalizedConstants.msgYes,
+            );
+            if (confirm !== LocalizedConstants.msgYes) {
+                return;
+            }
+        }
+
+        const rowIdToSelectionMap = new Map<number, ISlickRange[]>();
+        const rowIdToRowMap = new Map<number, DbCellValue[]>();
+
+        // Fetch all ranges in parallel; fill the maps as results come in
+        await Promise.all(
+            selection.map(async (range) => {
+                const count = range.toRow - range.fromRow + 1;
+                const result = await this.getRows(range.fromRow, count, batchId, resultId);
+                this.getRowMappings(
+                    result.resultSubset.rows,
+                    range,
+                    rowIdToSelectionMap,
+                    rowIdToRowMap,
+                );
+            }),
+        );
+
+        return calculateSelectionSummaryFromData(rowIdToRowMap, rowIdToSelectionMap);
+    }
+
     public async toggleSqlCmd(): Promise<boolean> {
         const queryExecuteOptions: QueryExecutionOptions = { options: {} };
         queryExecuteOptions.options["isSqlCmdMode"] = !this.isSqlCmd;

src/models/sqlOutputContentProvider.ts

@@ -242,6 +242,17 @@ export class SqlOutputContentProvider {
             .queryRunner.copyResultsAsInsertInto(selection, batchId, resultId);
     }
 
+    public generateSelectionSummaryData(
+        uri: string,
+        batchId: number,
+        resultId: number,
+        selection: Interfaces.ISlickRange[],
+    ): Promise<qr.SelectionSummaryStats> {
+        return this._queryResultsMap
+            .get(uri)
+            .queryRunner.generateSelectionSummaryData(selection, batchId, resultId);
+    }
+
     public editorSelectionRequestHandler(uri: string, selection: ISelectionData): void {
         void this._queryResultsMap.get(uri).queryRunner.setEditorSelection(selection);
     }

src/queryResult/utils.ts

@@ -20,6 +20,7 @@ import { QueryResultWebviewController } from "./queryResultWebViewController";
 import store, { SubKeys } from "./singletonStore";
 import { JsonFormattingEditProvider } from "../utils/jsonFormatter";
 import * as LocalizedConstants from "../constants/locConstants";
+import { generateGuid } from "../models/utils";
 
 export const MAX_VIEW_COLUMN = 9;
 
@@ -312,13 +313,34 @@ export function registerCommonRequestHandlers(
         return store.get(message.uri, SubKeys.PaneScrollPosition) ?? { scrollTop: 0 };
     });
 
-    webviewController.onRequest(qr.SetSelectionSummaryRequest.type, async (message) => {
+    let lastSummaryRequestId = "";
+    webviewController.onNotification(qr.SetSelectionSummaryRequest.type, async (message) => {
+        let currentId = generateGuid();
+        lastSummaryRequestId = currentId;
+
+        // Fetch all the data needed for the summary
+        const summaryData = await webviewViewController
+            .getSqlOutputContentProvider()
+            .generateSelectionSummaryData(
+                message.uri,
+                message.batchId,
+                message.resultId,
+                message.selection,
+            );
+
         const controller =
             webviewController instanceof QueryResultWebviewPanelController
                 ? webviewController.getQueryResultWebviewViewController()
                 : webviewController;
 
-        controller.updateSelectionSummaryStatusItem(message.summary);
+        /**
+         * Only update the summary if this is the latest request. This ensures that if multiple
+         * requests are made in quick succession, we don't overwrite the summary with stale data.
+         */
+        if (lastSummaryRequestId !== currentId) {
+            return;
+        }
+        controller.updateSelectionSummaryStatusItem(summaryData);
     });
 
     webviewController.registerReducer("setResultTab", async (state, payload) => {
@@ -414,3 +436,182 @@ export function countResultSets(
     }
     return count;
 }
+
+/**
+ * Calculate selection summary statistics for grid selections
+ * @param selections Array of grid selection ranges
+ * @param grid Mock grid object with getCellNode method
+ * @param isSelection Whether this is an actual selection (true) or clearing stats (false)
+ * @returns Promise<SelectionSummaryStats> Summary statistics
+ */
+export async function selectionSummaryHelper(
+    selections: qr.ISlickRange[],
+    grid: {
+        getCellNode: (row: number, col: number) => HTMLElement | undefined;
+        getColumns: () => any[];
+    },
+    isSelection: boolean,
+): Promise<qr.SelectionSummaryStats> {
+    const summary: qr.SelectionSummaryStats = {
+        count: -1,
+        average: "",
+        sum: 0,
+        min: 0,
+        max: 0,
+        removeSelectionStats: !isSelection,
+        distinctCount: -1,
+        nullCount: -1,
+    };
+
+    if (!isSelection) {
+        return summary;
+    }
+
+    const columns = grid.getColumns();
+    if (!columns || columns.length === 0) {
+        return summary;
+    }
+
+    if (!selections || selections.length === 0) {
+        return summary;
+    }
+
+    // Reset values for actual calculation
+    summary.count = 0;
+    summary.distinctCount = 0;
+    summary.nullCount = 0;
+    summary.min = Infinity;
+    summary.max = -Infinity;
+    summary.removeSelectionStats = false;
+
+    const distinct = new Set<string>();
+    let numericCount = 0;
+
+    const isFiniteNumber = (v: string): boolean => {
+        const n = Number(v);
+        return Number.isFinite(n);
+    };
+
+    for (const selection of selections) {
+        for (let row = selection.fromRow; row <= selection.toRow; row++) {
+            for (let col = selection.fromCell; col <= selection.toCell; col++) {
+                const cell = grid.getCellNode(row, col);
+                if (!cell) {
+                    continue;
+                }
+
+                summary.count++;
+                const cellText = cell.innerText;
+
+                if (cellText === "NULL" || cellText === null || cellText === undefined) {
+                    summary.nullCount++;
+                    continue;
+                }
+
+                distinct.add(cellText);
+
+                if (isFiniteNumber(cellText)) {
+                    const n = Number(cellText);
+                    numericCount++;
+                    summary.sum += n;
+                    if (n < summary.min) summary.min = n;
+                    if (n > summary.max) summary.max = n;
+                }
+            }
+        }
+    }
+
+    summary.distinctCount = distinct.size;
+
+    // Only compute average when we actually saw numeric cells
+    if (numericCount > 0) {
+        summary.average = (summary.sum / numericCount).toFixed(3);
+    } else {
+        summary.average = "";
+    }
+
+    // Normalize min/max if there were no numeric values
+    if (!Number.isFinite(summary.min)) summary.min = 0;
+    if (!Number.isFinite(summary.max)) summary.max = 0;
+
+    return summary;
+}
+
+/**
+ * Calculate selection summary statistics from database row data
+ * @param rowIdToRowMap Map of row IDs to database cell value arrays
+ * @param rowIdToSelectionMap Map of row IDs to selection ranges
+ * @returns SelectionSummaryStats Summary statistics
+ */
+export function calculateSelectionSummaryFromData(
+    rowIdToRowMap: Map<number, qr.DbCellValue[]>,
+    rowIdToSelectionMap: Map<number, qr.ISlickRange[]>,
+): qr.SelectionSummaryStats {
+    const summary: qr.SelectionSummaryStats = {
+        count: 0,
+        average: "",
+        sum: 0,
+        min: Infinity,
+        max: -Infinity,
+        removeSelectionStats: false,
+        distinctCount: 0,
+        nullCount: 0,
+    };
+
+    if (rowIdToRowMap.size === 0) {
+        // No data; normalize min/max to 0 to match prior behavior
+        summary.min = 0;
+        summary.max = 0;
+        return summary;
+    }
+
+    const distinct = new Set<string>();
+    let numericCount = 0;
+
+    const isFiniteNumber = (v: string): boolean => {
+        const n = Number(v);
+        return Number.isFinite(n);
+    };
+
+    for (const [rowId, row] of rowIdToRowMap) {
+        const rowSelections = rowIdToSelectionMap.get(rowId) ?? [];
+        for (const sel of rowSelections) {
+            const start = Math.max(0, sel.fromCell);
+            const end = Math.min(row.length - 1, sel.toCell);
+            for (let c = start; c <= end; c++) {
+                const cell = row[c];
+                summary.count++;
+
+                if (cell?.isNull) {
+                    summary.nullCount++;
+                    continue;
+                }
+
+                const display = cell?.displayValue ?? "";
+                distinct.add(display);
+
+                if (isFiniteNumber(display)) {
+                    const n = Number(display);
+                    numericCount++;
+                    summary.sum += n;
+                    if (n < summary.min) summary.min = n;
+                    if (n > summary.max) summary.max = n;
+                } else {
+                    // There is at least one non-numeric (non-null) value
+                    summary.removeSelectionStats = false;
+                }
+            }
+        }
+    }
+
+    summary.distinctCount = distinct.size;
+
+    // Only compute average when we actually saw numeric cells and round to 2 decimal places
+    summary.average = numericCount > 0 ? (summary.sum / numericCount).toFixed(2) : "";
+
+    // Normalize min/max if there were no numeric values
+    if (!Number.isFinite(summary.min)) summary.min = 0;
+    if (!Number.isFinite(summary.max)) summary.max = 0;
+
+    return summary;
+}

src/reactviews/common/keys.ts

@@ -1,14 +1,16 @@
-/*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the MIT License. See License.txt in the project root for license information.
- *--------------------------------------------------------------------------------------------*/
-
-export enum Keys {
-    ArrowDown = "ArrowDown",
-    ArrowUp = "ArrowUp",
-    Enter = "Enter",
-    Escape = "Escape",
-    Space = " ",
-    c = "c",
-    a = "a",
-}
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export enum Keys {
+    Enter = "Enter",
+    Escape = "Escape",
+    ArrowLeft = "ArrowLeft",
+    ArrowRight = "ArrowRight",
+    ArrowUp = "ArrowUp",
+    ArrowDown = "ArrowDown",
+    Space = " ",
+    c = "c",
+    a = "a",
+}