Show warning when SQL notebook result set is incomplete (#21506)

* Show error messages with partial results

* Generated loc changes

* Add unit tests

Author: lewis-sanchez

extensions/mssql/l10n/bundle.l10n.json

@@ -2002,6 +2002,13 @@
     },
     "(Command completed successfully)": "(Command completed successfully)",
     "(0 rows)": "(0 rows)",
+    "Warning: Result set is incomplete. Showing {0} of {1} rows. The full result set could not be loaded./{0} is the number of rows actually returned{1} is the total number of rows expected": {
+        "message": "Warning: Result set is incomplete. Showing {0} of {1} rows. The full result set could not be loaded.",
+        "comment": [
+            "{0} is the number of rows actually returned",
+            "{1} is the total number of rows expected"
+        ]
+    },
     "({0} row)/{0} is the number of rows (singular)": {
         "message": "({0} row)",
         "comment": ["{0} is the number of rows (singular)"]

extensions/mssql/src/constants/locConstants.ts

@@ -659,6 +659,17 @@ export class Notebooks {
     }
     public static commandCompletedSuccessfully = l10n.t("(Command completed successfully)");
     public static zeroRows = l10n.t("(0 rows)");
+    public static resultSetTruncated(actual: number, expected: number) {
+        return l10n.t({
+            message:
+                "Warning: Result set is incomplete. Showing {0} of {1} rows. The full result set could not be loaded.",
+            args: [actual, expected],
+            comment: [
+                "{0} is the number of rows actually returned",
+                "{1} is the total number of rows expected",
+            ],
+        });
+    }
     public static rowCountPlain(count: number) {
         if (count === 1) {
             return l10n.t({

extensions/mssql/src/notebooks/sqlNotebookController.ts

@@ -519,6 +519,20 @@ export class SqlNotebookController implements vscode.Disposable {
                     continue;
                 }
 
+                if (rs.rows.length < rs.rowCount) {
+                    outputs.push(
+                        new vscode.NotebookCellOutput([
+                            vscode.NotebookCellOutputItem.text(
+                                LocalizedConstants.Notebooks.resultSetTruncated(
+                                    rs.rows.length,
+                                    rs.rowCount,
+                                ),
+                                MIME_TEXT_PLAIN,
+                            ),
+                        ]),
+                    );
+                }
+
                 const plain = formatter.toPlain(rs.columnInfo, rs.rows);
                 outputs.push(
                     new vscode.NotebookCellOutput([

extensions/mssql/test/unit/notebooks/sqlNotebookController.test.ts

@@ -369,6 +369,88 @@ suite("SqlNotebookController", () => {
 
             expect(mockExecution.end).to.have.been.calledWith(false, sinon.match.number);
         });
+
+        test("shows truncation warning when result set is incomplete", async () => {
+            mockNotebookConnMgr.executeQueryString.resolves(
+                makeQueryResult({
+                    batches: [
+                        {
+                            batchSummary: makeBatchSummary(),
+                            messages: [],
+                            resultSets: [
+                                {
+                                    columnInfo: [makeColumn("id", "int")],
+                                    rows: [
+                                        [{ displayValue: "1", isNull: false }],
+                                        [{ displayValue: "2", isNull: false }],
+                                    ],
+                                    rowCount: 1000, // More rows exist than were returned
+                                },
+                            ],
+                            hasError: false,
+                        },
+                    ],
+                }),
+            );
+
+            const notebook = makeNotebook([{ text: "SELECT * FROM LargeTable" }]);
+            const cells = notebook.getCells();
+
+            await mockController.executeHandler(cells, notebook, mockController);
+
+            expect(mockExecution.replaceOutput).to.have.been.calledOnce;
+            const outputs = mockExecution.replaceOutput.firstCall.args[0];
+            expect(outputs).to.have.lengthOf(2);
+
+            // First output should be the truncation warning
+            const warningOutput = outputs[0];
+            expect(warningOutput.items[0].mime).to.equal("text/plain");
+            const warningText = new TextDecoder().decode(warningOutput.items[0].data);
+            expect(warningText).to.include("Warning: Result set is incomplete");
+            expect(warningText).to.include("2"); // Actual rows returned
+            expect(warningText).to.include("1000"); // Total rows available
+
+            expect(mockExecution.end).to.have.been.calledWith(true, sinon.match.number);
+        });
+
+        test("does not show truncation warning when result set is complete", async () => {
+            mockNotebookConnMgr.executeQueryString.resolves(
+                makeQueryResult({
+                    batches: [
+                        {
+                            batchSummary: makeBatchSummary(),
+                            messages: [],
+                            resultSets: [
+                                {
+                                    columnInfo: [makeColumn("id", "int")],
+                                    rows: [
+                                        [{ displayValue: "1", isNull: false }],
+                                        [{ displayValue: "2", isNull: false }],
+                                    ],
+                                    rowCount: 2, // Counts match - no truncation
+                                },
+                            ],
+                            hasError: false,
+                        },
+                    ],
+                }),
+            );
+
+            const notebook = makeNotebook([{ text: "SELECT * FROM SmallTable" }]);
+            const cells = notebook.getCells();
+
+            await mockController.executeHandler(cells, notebook, mockController);
+
+            expect(mockExecution.replaceOutput).to.have.been.calledOnce;
+            const outputs = mockExecution.replaceOutput.firstCall.args[0];
+            expect(outputs).to.have.lengthOf(1); // Only the result set output, no warning
+
+            // The output should be the result set, not a warning
+            const resultOutput = outputs[0];
+            expect(resultOutput.items[0].mime).to.equal("application/vnd.mssql.query-result");
+
+            expect(mockExecution.end).to.have.been.calledWith(true, sinon.match.number);
+        });
     });
 
     suite("executeCell — magic commands", () => {

localization/xliff/vscode-mssql.xlf

@@ -6393,6 +6393,11 @@
     <trans-unit id="++CODE++e981ddae45d8f4ca53f1ccbe613ad254a041dacf65a06026099a6302d332113b">
       <source xml:lang="en">Warning</source>
     </trans-unit>
+    <trans-unit id="++CODE++4dd8d9584066dcf1db567b8f698595cd5da2385469531e7c7684fa0792be1fcc">
+      <source xml:lang="en">Warning: Result set is incomplete. Showing {0} of {1} rows. The full result set could not be loaded.</source>
+      <note>{0} is the number of rows actually returned
+{1} is the total number of rows expected</note>
+    </trans-unit>
     <trans-unit id="++CODE++4d7371c8931fdd0178f7cbdabd1fd8260039a824a5c0912e880840b4d841f41e">
       <source xml:lang="en">Warnings detected. Please review the changes.</source>
     </trans-unit>