Execute SQL statement at cursor location (#934)

* Execute current statement at cursor location

* Refactor to add explicti runStatement methods

* Fix a couple bugs

* Bump SQL Tools Service to latest build

* Add null check to fix test breaks

* Fix another test break

package.json

@@ -37,6 +37,7 @@
     "onLanguage:sql",
     "onCommand:extension.connect",
     "onCommand:extension.runQuery",
+    "onCommand:extension.runCurrentStatement",
     "onCommand:extension.disconnect",
     "onCommand:extension.manageProfiles",
     "onCommand:extension.chooseDatabase",
@@ -158,6 +159,11 @@
         "title": "Execute Query",
         "category": "MS SQL"
       },
+      {
+        "command": "extension.runCurrentStatement",
+        "title": "Execute Current Statement",
+        "category": "MS SQL"
+      },
       {
         "command": "extension.cancelQuery",
         "title": "Cancel Query",

src/configurations/dev.config.json

@@ -1,7 +1,7 @@
 {
     "service": {
         "downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/v{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-        "version": "1.1.0-alpha.3",
+        "version": "1.1.0-alpha.5",
         "downloadFileNames": {
             "Windows_7_86": "win-x86-netcoreapp2.0.zip",
             "Windows_7_64": "win-x64-netcoreapp2.0.zip",

src/configurations/production.config.json

@@ -1,7 +1,7 @@
 {
     "service": {
         "downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/v{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-        "version": "1.1.0-alpha.3",
+        "version": "1.1.0-alpha.5",
         "downloadFileNames": {
             "Windows_7_86": "win-x86-netcoreapp2.0.zip",
             "Windows_7_64": "win-x64-netcoreapp2.0.zip",

src/constants/constants.ts

@@ -7,6 +7,7 @@ export const outputChannelName = 'MSSQL';
 export const connectionConfigFilename = 'settings.json';
 export const connectionsArrayName = 'mssql.connections';
 export const cmdRunQuery = 'extension.runQuery';
+export const cmdRunCurrentStatement = 'extension.runCurrentStatement';
 export const cmdCancelQuery = 'extension.cancelQuery';
 export const cmdConnect = 'extension.connect';
 export const cmdDisconnect = 'extension.disconnect';

src/controllers/mainController.ts

@@ -9,7 +9,7 @@ import vscode = require('vscode');
 import Constants = require('../constants/constants');
 import LocalizedConstants = require('../constants/localizedConstants');
 import Utils = require('../models/utils');
-import { SqlOutputContentProvider } from '../models/SqlOutputContentProvider';
+import { SqlOutputContentProvider } from '../models/sqlOutputContentProvider';
 import { RebuildIntelliSenseNotification } from '../models/contracts/languageService';
 import StatusView from '../views/statusView';
 import ConnectionManager from './connectionManager';
@@ -101,6 +101,8 @@ export default class MainController implements vscode.Disposable {
         this.registerCommand(Constants.cmdRunQuery);
         this._event.on(Constants.cmdRunQuery, () => { self.onRunQuery(); });
         this.registerCommand(Constants.cmdManageConnectionProfiles);
+        this._event.on(Constants.cmdRunCurrentStatement, () => { self.onRunCurrentStatement(); });
+        this.registerCommand(Constants.cmdRunCurrentStatement);
         this._event.on(Constants.cmdManageConnectionProfiles, () => { self.runAndLogErrors(self.onManageProfiles(), 'onManageProfiles'); });
         this.registerCommand(Constants.cmdChooseDatabase);
         this._event.on(Constants.cmdChooseDatabase, () => { self.runAndLogErrors(self.onChooseDatabase(), 'onChooseDatabase') ; } );
@@ -254,6 +256,42 @@ export default class MainController implements vscode.Disposable {
         }
     }
 
+    /**
+     * execute the SQL statement for the current cursor position
+     */
+    public onRunCurrentStatement(): void {
+        try {
+            if (!this.CanRunCommand()) {
+                return;
+            }
+
+            // check if we're connected and editing a SQL file
+            if (this.isRetryRequiredBeforeQuery(this.onRunCurrentStatement)) {
+                return;
+            }
+
+            let editor = this._vscodeWrapper.activeTextEditor;
+            let uri = this._vscodeWrapper.activeTextEditorUri;
+            let title = path.basename(editor.document.fileName);
+            let querySelection: ISelectionData;
+
+            Telemetry.sendTelemetryEvent('RunCurrentStatement');
+
+            // only the start line and column are used to determine the current statement
+            querySelection = {
+                startLine: editor.selection.start.line,
+                startColumn: editor.selection.start.character,
+                endLine: 0,
+                endColumn: 0
+            };
+
+            this._outputContentProvider.runCurrentStatement(this._statusview, uri, querySelection, title);
+        } catch (err) {
+            Telemetry.sendTelemetryEventForException(err, 'onRunCurrentStatement');
+        }
+
+    }
+
     /**
      * get the T-SQL query from the editor, run it and show output
      */
@@ -262,57 +300,75 @@ export default class MainController implements vscode.Disposable {
             if (!this.CanRunCommand()) {
                 return;
             }
-            const self = this;
-            if (!this._vscodeWrapper.isEditingSqlFile) {
-                // Prompt the user to change the language mode to SQL before running a query
-                this._connectionMgr.connectionUI.promptToChangeLanguageMode().then( result => {
-                    if (result) {
-                        self.onRunQuery();
-                    }
-                }).catch(err => {
-                    self._vscodeWrapper.showErrorMessage(LocalizedConstants.msgError + err);
-                });
-            } else if (!this._connectionMgr.isConnected(this._vscodeWrapper.activeTextEditorUri)) {
-                // If we are disconnected, prompt the user to choose a connection before executing
-                this.onNewConnection().then(result => {
-                    if (result) {
-                        self.onRunQuery();
-                    }
-                }).catch(err => {
-                    self._vscodeWrapper.showErrorMessage(LocalizedConstants.msgError + err);
-                });
-            } else {
-                let editor = this._vscodeWrapper.activeTextEditor;
-                let uri = this._vscodeWrapper.activeTextEditorUri;
-                let title = path.basename(editor.document.fileName);
-                let querySelection: ISelectionData;
 
-                // Calculate the selection if we have a selection, otherwise we'll use null to indicate
-                // the entire document is the selection
-                if (!editor.selection.isEmpty) {
-                    let selection = editor.selection;
-                    querySelection = {
-                        startLine: selection.start.line,
-                        startColumn: selection.start.character,
-                        endLine: selection.end.line,
-                        endColumn: selection.end.character
-                    };
-                }
-
-                // Trim down the selection. If it is empty after selecting, then we don't execute
-                let selectionToTrim = editor.selection.isEmpty ? undefined : editor.selection;
-                if (editor.document.getText(selectionToTrim).trim().length === 0) {
-                    return;
-                }
-
-                Telemetry.sendTelemetryEvent('RunQuery');
-
-                this._outputContentProvider.runQuery(this._statusview, uri, querySelection, title);
+            // check if we're connected and editing a SQL file
+            if (this.isRetryRequiredBeforeQuery(this.onRunQuery)) {
+                return;
             }
-        } catch (err) {
-            Telemetry.sendTelemetryEventForException(err, 'OnRunquery');
-        }
 
+            let editor = this._vscodeWrapper.activeTextEditor;
+            let uri = this._vscodeWrapper.activeTextEditorUri;
+            let title = path.basename(editor.document.fileName);
+            let querySelection: ISelectionData;
+
+            // Calculate the selection if we have a selection, otherwise we'll use null to indicate
+            // the entire document is the selection
+            if (!editor.selection.isEmpty) {
+                let selection = editor.selection;
+                querySelection = {
+                    startLine: selection.start.line,
+                    startColumn: selection.start.character,
+                    endLine: selection.end.line,
+                    endColumn: selection.end.character
+                };
+            }
+
+            // Trim down the selection. If it is empty after selecting, then we don't execute
+            let selectionToTrim = editor.selection.isEmpty ? undefined : editor.selection;
+            if (editor.document.getText(selectionToTrim).trim().length === 0) {
+                return;
+            }
+
+            Telemetry.sendTelemetryEvent('RunQuery');
+
+            this._outputContentProvider.runQuery(this._statusview, uri, querySelection, title);
+        } catch (err) {
+            Telemetry.sendTelemetryEventForException(err, 'onRunQuery');
+        }
+    }
+
+    /**
+     * Check if the state is ready to execute a query and retry
+     * the query execution method if needed
+     */
+    public isRetryRequiredBeforeQuery(retryMethod: any): boolean {
+        const self = this;
+        if (!this._vscodeWrapper.isEditingSqlFile) {
+            // Prompt the user to change the language mode to SQL before running a query
+            this._connectionMgr.connectionUI.promptToChangeLanguageMode().then( result => {
+                if (result) {
+                    retryMethod();
+                }
+            }).catch(err => {
+                self._vscodeWrapper.showErrorMessage(LocalizedConstants.msgError + err);
+            });
+            return true;
+
+        } else if (!this._connectionMgr.isConnected(this._vscodeWrapper.activeTextEditorUri)) {
+            // If we are disconnected, prompt the user to choose a connection before executing
+            this.onNewConnection().then(result => {
+                if (result) {
+                    retryMethod();
+                }
+            }).catch(err => {
+                self._vscodeWrapper.showErrorMessage(LocalizedConstants.msgError + err);
+            });
+            return true;
+        } else {
+
+            // we don't need to do anything to configure environment before running query
+            return false;
+        }
     }
 
     /**

src/controllers/queryNotificationHandler.ts

@@ -2,7 +2,7 @@
 *  Class for handler and distributing notification coming from the
 *  service layer
 */
-import QueryRunner from './QueryRunner';
+import QueryRunner from './queryRunner';
 import SqlToolsServiceClient from '../languageservice/serviceclient';
 import {
     QueryExecuteCompleteNotification,

src/controllers/queryRunner.ts

@@ -3,16 +3,17 @@ import { EventEmitter } from 'events';
 
 import StatusView from '../views/statusView';
 import SqlToolsServerClient from '../languageservice/serviceclient';
-import {QueryNotificationHandler} from './QueryNotificationHandler';
+import {QueryNotificationHandler} from './queryNotificationHandler';
 import VscodeWrapper from './vscodeWrapper';
 import { BatchSummary, QueryExecuteParams, QueryExecuteRequest,
+    QueryExecuteStatementParams, QueryExecuteStatementRequest,
     QueryExecuteCompleteNotificationResult, QueryExecuteSubsetResult,
     QueryExecuteResultSetCompleteNotificationParams,
     QueryExecuteSubsetParams, QueryExecuteSubsetRequest,
     QueryExecuteMessageParams,
     QueryExecuteBatchNotificationParams } from '../models/contracts/queryExecute';
-import { QueryDisposeParams, QueryDisposeRequest } from '../models/contracts/QueryDispose';
-import { QueryCancelParams, QueryCancelResult, QueryCancelRequest } from '../models/contracts/QueryCancel';
+import { QueryDisposeParams, QueryDisposeRequest } from '../models/contracts/queryDispose';
+import { QueryCancelParams, QueryCancelResult, QueryCancelRequest } from '../models/contracts/queryCancel';
 import { ISlickRange, ISelectionData } from '../models/interfaces';
 import Constants = require('../constants/constants');
 import LocalizedConstants = require('../constants/localizedConstants');
@@ -110,35 +111,63 @@ export default class QueryRunner {
         return this._client.sendRequest(QueryCancelRequest.type, cancelParams);
     }
 
+    // Pulls the query text from the current document/selection and initiates the query
+    public runStatement(line: number, column: number): Thenable<void> {
+        return this.doRunQuery(
+            <ISelectionData>{ startLine: line, startColumn: column, endLine: 0, endColumn: 0 },
+            (onSuccess, onError) => {
+                // Put together the request
+                let queryDetails: QueryExecuteStatementParams = {
+                    ownerUri: this._uri,
+                    line: line,
+                    column: column
+                };
+
+                // Send the request to execute the query
+                return this._client.sendRequest(QueryExecuteStatementRequest.type, queryDetails).then(onSuccess, onError);
+            });
+    }
+
     // Pulls the query text from the current document/selection and initiates the query
     public runQuery(selection: ISelectionData): Thenable<void> {
+        return this.doRunQuery(
+            selection,
+            (onSuccess, onError) => {
+               // Put together the request
+                let queryDetails: QueryExecuteParams = {
+                    ownerUri: this._uri,
+                    querySelection: selection
+                };
+
+                // Send the request to execute the query
+                return this._client.sendRequest(QueryExecuteRequest.type, queryDetails).then(onSuccess, onError);
+            });
+    }
+
+    // Pulls the query text from the current document/selection and initiates the query
+    private doRunQuery(selection: ISelectionData, queryCallback: any): Thenable<void> {
         const self = this;
         this._vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgStartedExecute, this._uri));
 
-        // Put together the request
-        let queryDetails: QueryExecuteParams = {
-            ownerUri: this._uri,
-            querySelection: selection
-        };
-
         // Update internal state to show that we're executing the query
         this._resultLineOffset = selection ? selection.startLine : 0;
         this._isExecuting = true;
         this._totalElapsedMilliseconds = 0;
         this._statusView.executingQuery(this.uri);
 
-        // Send the request to execute the query
-        return this._client.sendRequest(QueryExecuteRequest.type, queryDetails).then(result => {
+        let onSuccess = (result) => {
             // The query has started, so lets fire up the result pane
             self.eventEmitter.emit('start');
-            self._notificationHandler.registerRunner(self, queryDetails.ownerUri);
-        }, error => {
-            // Attempting to launch the query failed, show the error message
+            self._notificationHandler.registerRunner(self, self._uri);
+        };
+        let onError = (error) => {
             self._statusView.executedQuery(self.uri);
             self._isExecuting = false;
             // TODO: localize
             self._vscodeWrapper.showErrorMessage('Execution failed: ' + error.message);
-        });
+        };
+
+        return queryCallback(onSuccess, onError);
     }
 
     // handle the result of the notification

src/models/contracts/queryExecute.ts

@@ -95,11 +95,25 @@ export namespace QueryExecuteRequest {
     };
 }
 
+export namespace QueryExecuteStatementRequest {
+    export const type: RequestType<QueryExecuteStatementParams, QueryExecuteResult, void> = {
+        get method(): string {
+            return 'query/executedocumentstatement';
+        }
+    };
+}
+
 export class QueryExecuteParams {
     ownerUri: string;
     querySelection: ISelectionData;
 }
 
+export class QueryExecuteStatementParams {
+    ownerUri: string;
+    line: number;
+    column: number;
+}
+
 export class QueryExecuteResult {}
 
 // --------------------------------- < Query Results Request > ------------------------------------------

src/models/sqlOutputContentProvider.ts

@@ -6,7 +6,7 @@ import LocalizedConstants = require('../constants/localizedConstants');
 import LocalWebService from '../controllers/localWebService';
 import Utils = require('./utils');
 import Interfaces = require('./interfaces');
-import QueryRunner from '../controllers/QueryRunner';
+import QueryRunner from '../controllers/queryRunner';
 import ResultsSerializer from  '../models/resultsSerializer';
 import StatusView from '../views/statusView';
 import VscodeWrapper from './../controllers/vscodeWrapper';
@@ -218,9 +218,49 @@ export class SqlOutputContentProvider implements vscode.TextDocumentContentProvi
             : this._queryResultsMap.get(uri).queryRunner.isExecutingQuery;
     }
 
-    public runQuery(statusView, uri: string, selection: ISelectionData, title: string): void {
-        // Reuse existing query runner if it exists
+    public runQuery(
+            statusView: any, uri: string,
+            selection: ISelectionData, title: string): void {
+        // execute the query with a query runner
+        this.runQueryCallback(statusView, uri, selection, title,
+            (queryRunner) => {
+                if (queryRunner) {
+                    queryRunner.runQuery(selection);
+                }
+            });
+    }
+
+    public runCurrentStatement(
+            statusView: any, uri: string,
+            selection: ISelectionData, title: string): void {
+        // execute the statement with a query runner
+        this.runQueryCallback(statusView, uri, selection, title,
+            (queryRunner) => {
+                if (queryRunner) {
+                    queryRunner.runStatement(selection.startLine, selection.startColumn);
+                }
+            });
+    }
+
+    private runQueryCallback(
+            statusView: any, uri: string,
+            selection: ISelectionData, title: string,
+            queryCallback: any): void {
+
         let resultsUri = this.getResultsUri(uri);
+        let queryRunner = this.createQueryRunner(statusView, uri, resultsUri, selection, title);
+        if (queryRunner) {
+            queryCallback(queryRunner);
+
+            let paneTitle = Utils.formatString(LocalizedConstants.titleResultsPane, queryRunner.title);
+            // Always run this command even if just updating to avoid a bug - tfs 8686842
+            this.displayResultPane(resultsUri, paneTitle);
+        }
+    }
+
+    private createQueryRunner(statusView: any, uri: string, resultsUri: string,
+                              selection: ISelectionData, title: string): QueryRunner {
+        // Reuse existing query runner if it exists
         let queryRunner: QueryRunner;
 
         if (this._queryResultsMap.has(resultsUri)) {
@@ -280,10 +320,7 @@ export class SqlOutputContentProvider implements vscode.TextDocumentContentProvi
             this._queryResultsMap.set(resultsUri, new QueryRunnerState(queryRunner));
         }
 
-        queryRunner.runQuery(selection);
-        let paneTitle = Utils.formatString(LocalizedConstants.titleResultsPane, queryRunner.title);
-        // Always run this command even if just updating to avoid a bug - tfs 8686842
-        this.displayResultPane(resultsUri, paneTitle);
+        return queryRunner;
     }
 
     // Function to render resultspane content

test/queryNotificationHandler.test.ts

@@ -1,7 +1,7 @@
 import * as TypeMoq from 'typemoq';
 import assert = require('assert');
-import QueryRunner from './../src/controllers/QueryRunner';
-import { QueryNotificationHandler } from './../src/controllers/QueryNotificationHandler';
+import QueryRunner from './../src/controllers/queryRunner';
+import { QueryNotificationHandler } from './../src/controllers/queryNotificationHandler';
 import { NotificationHandler } from 'vscode-languageclient';
 
 // TESTS //////////////////////////////////////////////////////////////////////////////////////////

test/queryRunner.test.ts

@@ -1,9 +1,9 @@
 import * as TypeMoq from 'typemoq';
 import assert = require('assert');
 import { EventEmitter } from 'events';
-import QueryRunner from './../src/controllers/QueryRunner';
-import { QueryNotificationHandler } from './../src/controllers/QueryNotificationHandler';
-import { SqlOutputContentProvider } from './../src/models/SqlOutputContentProvider';
+import QueryRunner from './../src/controllers/queryRunner';
+import { QueryNotificationHandler } from './../src/controllers/queryNotificationHandler';
+import { SqlOutputContentProvider } from './../src/models/sqlOutputContentProvider';
 import SqlToolsServerClient from './../src/languageservice/serviceclient';
 import {
     QueryExecuteParams,
@@ -17,7 +17,7 @@ import VscodeWrapper from './../src/controllers/vscodeWrapper';
 import StatusView from './../src/views/statusView';
 import * as Constants from '../src/constants/constants';
 import * as QueryExecuteContracts from '../src/models/contracts/queryExecute';
-import * as QueryDisposeContracts from '../src/models/contracts/QueryDispose';
+import * as QueryDisposeContracts from '../src/models/contracts/queryDispose';
 import {
     ISlickRange,
     ISelectionData

test/sqlOutputContentProvider.test.ts

@@ -1,6 +1,6 @@
 'use strict';
 
-import { SqlOutputContentProvider } from '../src/models/SqlOutputContentProvider';
+import { SqlOutputContentProvider } from '../src/models/sqlOutputContentProvider';
 import VscodeWrapper from '../src/controllers/vscodeWrapper';
 import StatusView from '../src/views/statusView';
 import * as stubs from './stubs';

test/webServiceRequestHandlers.test.ts

@@ -1,8 +1,8 @@
 'use strict';
 
-import { SqlOutputContentProvider } from '../src/models/SqlOutputContentProvider';
+import { SqlOutputContentProvider } from '../src/models/sqlOutputContentProvider';
 import VscodeWrapper from '../src/controllers/vscodeWrapper';
-import QueryRunner from '../src/controllers/QueryRunner';
+import QueryRunner from '../src/controllers/queryRunner';
 import { ISelectionData } from '../src/models/interfaces';
 import { QueryExecuteSubsetResult, ResultSetSubset } from '../src/models/contracts/queryExecute';
 import StatusView from '../src/views/statusView';