Fix Execute Current Statement selection behavior

extensions/mssql/package.nls.json

@@ -1,6 +1,6 @@
 {
     "mssql.runQuery": "Execute Query (MSSQL)",
-    "mssql.runCurrentStatement": "Execute Current Statement (MSSQL)",
+    "mssql.runCurrentStatement": "Execute Selection or Current Statement (MSSQL)",
     "mssql.cancelQuery": "Cancel Query (MSSQL)",
     "mssql.revealQueryResult": "Reveal Query Result (MSSQL)",
     "mssql.toggleQueryResultPanel": "Toggle Query Result Panel Visibility (MSSQL)",

extensions/mssql/src/controllers/mainController.ts

@@ -2761,6 +2761,37 @@ export default class MainController implements vscode.Disposable {
                 return;
             }
 
+            // Do not execute when there are multiple selections in the editor until it can be properly handled.
+            // Otherwise only the first selection will be executed and cause unexpected issues.
+            if (editor.selections?.length > 1) {
+                self._vscodeWrapper.showErrorMessage(
+                    LocalizedConstants.msgMultipleSelectionModeNotSupported,
+                );
+                return;
+            }
+
+            if (!editor.selection.isEmpty) {
+                if (editor.document.getText(editor.selection).trim().length === 0) {
+                    return;
+                }
+
+                let selection = editor.selection;
+                let querySelection: ISelectionData = {
+                    startLine: selection.start.line,
+                    startColumn: selection.start.character,
+                    endLine: selection.end.line,
+                    endColumn: selection.end.character,
+                };
+
+                await self._outputContentProvider.runQuery(
+                    self._statusview,
+                    uri,
+                    querySelection,
+                    title,
+                );
+                return;
+            }
+
             // only the start line and column are used to determine the current statement
             let querySelection: ISelectionData = {
                 startLine: editor.selection.start.line,

extensions/mssql/test/unit/mainController.test.ts

@@ -29,9 +29,16 @@ type MainControllerTestAccess = {
     validateTextDocumentHasFocus(): boolean;
     _vscodeWrapper: {
         activeTextEditorUri?: string;
+        activeTextEditor?: vscode.TextEditor;
         isEditingSqlFile?: boolean;
         showInformationMessage(message: string): Thenable<unknown> | void;
+        showErrorMessage?(message: string): Thenable<unknown> | void;
     };
+    _outputContentProvider: {
+        runCurrentStatement: sinon.SinonStub;
+        runQuery: sinon.SinonStub;
+    };
+    _statusview: unknown;
     openCopilotChatFromUi(args?: CopilotChat.OpenFromUiArgs): Promise<void>;
     findChatOpenAgentCommand(): Promise<string | undefined>;
     registerLanguageModelTools(): void;
@@ -60,6 +67,24 @@ suite("MainController Tests", function () {
         } as unknown as vscode.TextEditor;
     }
 
+    function createQueryTextEditor(
+        selection: vscode.Selection,
+        fullText: string,
+        selectedText?: string,
+        selections: vscode.Selection[] = [selection],
+    ): vscode.TextEditor {
+        return {
+            document: {
+                uri: vscode.Uri.parse("file:///test/query.sql"),
+                fileName: "query.sql",
+                languageId: Constants.languageId,
+                getText: (range?: vscode.Range) => (range ? (selectedText ?? "") : fullText),
+            },
+            selection,
+            selections,
+        } as unknown as vscode.TextEditor;
+    }
+
     setup(() => {
         sandbox = sinon.createSandbox();
 
@@ -131,6 +156,119 @@ suite("MainController Tests", function () {
         expect(connectionManager.onManageProfiles).to.have.been.called;
     });
 
+    suite("onRunCurrentStatement Tests", () => {
+        let controllerAccess: MainControllerTestAccess;
+        let runCurrentStatementStub: sinon.SinonStub;
+        let runQueryStub: sinon.SinonStub;
+        let statusView: unknown;
+
+        setup(() => {
+            controllerAccess = accessMainController(mainController);
+            sandbox.stub(mainController, "ensureReadyToExecuteQuery").resolves(true);
+            sandbox.stub(vscodeWrapper, "activeTextEditorUri").get(() => "file:///test/query.sql");
+
+            runCurrentStatementStub = sandbox.stub().resolves();
+            runQueryStub = sandbox.stub().resolves();
+            controllerAccess._outputContentProvider = {
+                runCurrentStatement: runCurrentStatementStub,
+                runQuery: runQueryStub,
+            };
+            statusView = {};
+            controllerAccess._statusview = statusView;
+        });
+
+        test("runs the current statement when there is no selection", async () => {
+            const selection = new vscode.Selection(1, 7, 1, 7);
+            sandbox
+                .stub(vscodeWrapper, "activeTextEditor")
+                .get(() => createQueryTextEditor(selection, "select 'a';\nselect 'b';"));
+
+            await mainController.onRunCurrentStatement();
+
+            expect(runCurrentStatementStub).to.have.been.calledOnceWithExactly(
+                statusView,
+                "file:///test/query.sql",
+                {
+                    startLine: 1,
+                    startColumn: 7,
+                    endLine: 0,
+                    endColumn: 0,
+                },
+                "query.sql",
+            );
+            expect(runQueryStub).to.not.have.been.called;
+        });
+
+        test("runs selected text when there is a non-empty selection", async () => {
+            const selection = new vscode.Selection(0, 0, 0, 11);
+            sandbox
+                .stub(vscodeWrapper, "activeTextEditor")
+                .get(() =>
+                    createQueryTextEditor(selection, "select 'a'; select 'b';", "select 'a';"),
+                );
+
+            await mainController.onRunCurrentStatement();
+
+            expect(runQueryStub).to.have.been.calledOnceWithExactly(
+                statusView,
+                "file:///test/query.sql",
+                {
+                    startLine: 0,
+                    startColumn: 0,
+                    endLine: 0,
+                    endColumn: 11,
+                },
+                "query.sql",
+            );
+            expect(runCurrentStatementStub).to.not.have.been.called;
+        });
+
+        test("does not execute when the selection contains only whitespace", async () => {
+            const selection = new vscode.Selection(0, 0, 0, 4);
+            sandbox
+                .stub(vscodeWrapper, "activeTextEditor")
+                .get(() => createQueryTextEditor(selection, "    select 'a';", "    "));
+
+            await mainController.onRunCurrentStatement();
+
+            expect(runQueryStub).to.not.have.been.called;
+            expect(runCurrentStatementStub).to.not.have.been.called;
+        });
+
+        test("shows an error and does not execute when there are multiple selections", async () => {
+            const selection = new vscode.Selection(0, 0, 0, 11);
+            const secondSelection = new vscode.Selection(1, 0, 1, 11);
+            sandbox
+                .stub(vscodeWrapper, "activeTextEditor")
+                .get(() =>
+                    createQueryTextEditor(selection, "select 'a';\nselect 'b';", "select 'a';", [
+                        selection,
+                        secondSelection,
+                    ]),
+                );
+
+            await mainController.onRunCurrentStatement();
+
+            expect(vscodeWrapper.showErrorMessage).to.have.been.calledOnceWithExactly(
+                LocalizedConstants.msgMultipleSelectionModeNotSupported,
+            );
+            expect(runQueryStub).to.not.have.been.called;
+            expect(runCurrentStatementStub).to.not.have.been.called;
+        });
+
+        test("does not execute when the document is empty", async () => {
+            const selection = new vscode.Selection(0, 0, 0, 0);
+            sandbox
+                .stub(vscodeWrapper, "activeTextEditor")
+                .get(() => createQueryTextEditor(selection, ""));
+
+            await mainController.onRunCurrentStatement();
+
+            expect(runQueryStub).to.not.have.been.called;
+            expect(runCurrentStatementStub).to.not.have.been.called;
+        });
+    });
+
     test("Proxy settings are checked on initialization", async () => {
         const httpHelperWarnSpy = sandbox.spy(HttpClient.prototype, "warnOnInvalidProxySettings");
 

localization/xliff/vscode-mssql.xlf

@@ -7569,12 +7569,12 @@
     <trans-unit id="mssql.showEstimatedPlan">
       <source xml:lang="en">Estimated Plan (MSSQL)</source>
     </trans-unit>
-    <trans-unit id="mssql.runCurrentStatement">
-      <source xml:lang="en">Execute Current Statement (MSSQL)</source>
-    </trans-unit>
     <trans-unit id="mssql.runQuery">
       <source xml:lang="en">Execute Query (MSSQL)</source>
     </trans-unit>
+    <trans-unit id="mssql.runCurrentStatement">
+      <source xml:lang="en">Execute Selection or Current Statement (MSSQL)</source>
+    </trans-unit>
     <trans-unit id="mssql.copilot.explainQuery">
       <source xml:lang="en">Explain Query</source>
     </trans-unit>