Query Runner Tests (#10252)

* rework some code and write an inital test

* fix strict

* add more to standard test

* add to existing workflow test

* fix tests

* simplify the code

* add more tests

* remove bad import

* fix compile

* fix timestampiong

src/sql/workbench/services/query/common/gridDataProvider.ts

@@ -3,9 +3,9 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import * as azdata from 'azdata';
 import * as types from 'vs/base/common/types';
 import { SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
+import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
 
 export interface IGridDataProvider {
 
@@ -14,7 +14,7 @@ export interface IGridDataProvider {
 	 * @param rowStart 0-indexed start row to retrieve data from
 	 * @param numberOfRows total number of rows of data to retrieve
 	 */
-	getRowData(rowStart: number, numberOfRows: number): Thenable<azdata.QueryExecuteSubsetResult>;
+	getRowData(rowStart: number, numberOfRows: number): Thenable<ResultSetSubset>;
 
 	/**
 	 * Sends a copy request to copy data to the clipboard
@@ -65,8 +65,8 @@ export async function getResultsString(provider: IGridDataProvider, selection: S
 				}
 			}
 			// Iterate over the rows to paste into the copy string
-			for (let rowIndex: number = 0; rowIndex < result.resultSubset.rows.length; rowIndex++) {
-				let row = result.resultSubset.rows[rowIndex];
+			for (let rowIndex: number = 0; rowIndex < result.rows.length; rowIndex++) {
+				let row = result.rows[rowIndex];
 				let cellObjects = row.slice(range.fromCell, (range.toCell + 1));
 				// Remove newlines if requested
 				let cells = provider.shouldRemoveNewLines()

src/sql/workbench/services/query/common/query.ts

@@ -0,0 +1,69 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IRange } from 'vs/editor/common/core/range';
+
+export interface IColumn {
+	columnName: string;
+	isXml?: boolean;
+	isJson?: boolean;
+}
+
+export interface ResultSetSummary {
+	id: number;
+	batchId: number;
+	rowCount: number;
+	columnInfo: IColumn[];
+	complete: boolean;
+}
+
+export interface BatchStartSummary {
+	id: number;
+	executionStart: string;
+	range?: IRange;
+}
+
+export interface BatchSummary extends BatchStartSummary {
+	hasError: boolean;
+	resultSetSummaries: ResultSetSummary[];
+}
+
+export interface CompleteBatchSummary extends BatchSummary {
+	executionElapsed: string;
+	executionEnd: string;
+}
+
+export interface IQueryMessage {
+	batchId?: number;
+	isError: boolean;
+	time?: string;
+	message: string;
+	range?: IRange;
+}
+
+export interface IResultMessage {
+	batchId?: number;
+	isError: boolean;
+	time?: string;
+	message: string;
+}
+
+export interface QueryExecuteSubsetParams {
+	ownerUri: string;
+	batchIndex: number;
+	resultSetIndex: number;
+	rowsStartIndex: number;
+	rowsCount: number;
+}
+
+export interface ResultSetSubset {
+	rowCount: number;
+	rows: ICellValue[][];
+}
+
+export interface ICellValue {
+	displayValue: string;
+	isNull?: boolean;
+}

src/sql/workbench/services/query/common/queryManagement.ts

@@ -13,11 +13,24 @@ import { Event, Emitter } from 'vs/base/common/event';
 import { keys } from 'vs/base/common/map';
 import { assign } from 'vs/base/common/objects';
 import { IAdsTelemetryService, ITelemetryEventProperties } from 'sql/platform/telemetry/common/telemetry';
+import EditQueryRunner from 'sql/workbench/services/editData/common/editQueryRunner';
+import { IRange, Range } from 'vs/editor/common/core/range';
+import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
+import { isUndefined } from 'vs/base/common/types';
 
 export const SERVICE_ID = 'queryManagementService';
 
 export const IQueryManagementService = createDecorator<IQueryManagementService>(SERVICE_ID);
 
+export interface QueryCancelResult {
+	messages: string;
+}
+
+export interface ExecutionPlanOptions {
+	displayEstimatedQueryPlan?: boolean;
+	displayActualQueryPlan?: boolean;
+}
+
 export interface IQueryManagementService {
 	_serviceBrand: undefined;
 
@@ -28,13 +41,13 @@ export interface IQueryManagementService {
 	getRegisteredProviders(): string[];
 	registerRunner(runner: QueryRunner, uri: string): void;
 
-	cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult>;
-	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	cancelQuery(ownerUri: string): Promise<QueryCancelResult>;
+	runQuery(ownerUri: string, range: IRange, runOptions?: ExecutionPlanOptions): Promise<void>;
 	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
 	runQueryString(ownerUri: string, queryString: string): Promise<void>;
 	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
 	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult>;
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset>;
 	disposeQuery(ownerUri: string): Promise<void>;
 	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
 	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Promise<void>;
@@ -67,7 +80,7 @@ export interface IQueryManagementService {
  */
 export interface IQueryRequestHandler {
 	cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult>;
-	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: ExecutionPlanOptions): Promise<void>;
 	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
 	runQueryString(ownerUri: string, queryString: string): Promise<void>;
 	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
@@ -129,7 +142,7 @@ export class QueryManagementService implements IQueryManagementService {
 	// Handles logic to run the given handlerCallback at the appropriate time. If the given runner is
 	// undefined, the handlerCallback is put on the _handlerCallbackQueue to be run once the runner is set
 	// public for testing only
-	private enqueueOrRun(handlerCallback: (runnerParam: QueryRunner) => void, runner: QueryRunner): void {
+	private enqueueOrRun(handlerCallback: (runnerParam: QueryRunner) => void, runner?: QueryRunner): void {
 		if (runner === undefined) {
 			this._handlerCallbackQueue.push(handlerCallback);
 		} else {
@@ -137,9 +150,9 @@ export class QueryManagementService implements IQueryManagementService {
 		}
 	}
 
-	private _notify(ownerUri: string, sendNotification: (runner: QueryRunner) => void): void {
+	private _notify(ownerUri: string, sendNotification: (runner: QueryRunner | EditQueryRunner) => void): void {
 		let runner = this._queryRunners.get(ownerUri);
-		this.enqueueOrRun(sendNotification, runner!);
+		this.enqueueOrRun(sendNotification, runner);
 	}
 
 	public addQueryRequestHandler(queryType: string, handler: IQueryRequestHandler): IDisposable {
@@ -165,7 +178,7 @@ export class QueryManagementService implements IQueryManagementService {
 		return Array.from(keys(this._requestHandlers));
 	}
 
-	private addTelemetry(eventName: string, ownerUri: string, runOptions?: azdata.ExecutionPlanOptions): void {
+	private addTelemetry(eventName: string, ownerUri: string, runOptions?: ExecutionPlanOptions): void {
 		const providerId: string = this._connectionService.getProviderIdFromUri(ownerUri);
 		const data: ITelemetryEventProperties = {
 			provider: providerId,
@@ -197,51 +210,59 @@ export class QueryManagementService implements IQueryManagementService {
 		}
 	}
 
-	public cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult> {
+	public cancelQuery(ownerUri: string): Promise<QueryCancelResult> {
 		this.addTelemetry(TelemetryKeys.CancelQuery, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.cancelQuery(ownerUri);
 		});
 	}
-	public runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
+
+	public runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQuery, ownerUri, runOptions);
 		return this._runAction(ownerUri, (runner) => {
-			return runner.runQuery(ownerUri, selection, runOptions);
+			return runner.runQuery(ownerUri, rangeToSelectionData(range), runOptions);
 		});
 	}
+
 	public runQueryStatement(ownerUri: string, line: number, column: number): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQueryStatement, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
-			return runner.runQueryStatement(ownerUri, line, column);
+			return runner.runQueryStatement(ownerUri, line - 1, column - 1); // we are taking in a vscode IRange which is 1 indexed, but our api expected a 0 index
 		});
 	}
+
 	public runQueryString(ownerUri: string, queryString: string): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQueryString, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryString(ownerUri, queryString);
 		});
 	}
+
 	public runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryAndReturn(ownerUri, queryString);
 		});
 	}
+
 	public parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.parseSyntax(ownerUri, query);
 		});
 	}
-	public getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult> {
+
+	public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
 		return this._runAction(rowData.ownerUri, (runner) => {
-			return runner.getQueryRows(rowData);
+			return runner.getQueryRows(rowData).then(r => r.resultSubset);
 		});
 	}
+
 	public disposeQuery(ownerUri: string): Promise<void> {
 		this._queryRunners.delete(ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.disposeQuery(ownerUri);
 		});
 	}
+
 	public setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.setQueryExecutionOptions(ownerUri, options);
@@ -256,37 +277,38 @@ export class QueryManagementService implements IQueryManagementService {
 
 	public onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void {
 		this._notify(result.ownerUri, (runner: QueryRunner) => {
-			runner.handleQueryComplete(result);
+			runner.handleQueryComplete(result.batchSummaries.map(s => ({ ...s, range: selectionDataToRange(s.selection) })));
 		});
 	}
+
 	public onBatchStart(batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
 		this._notify(batchInfo.ownerUri, (runner: QueryRunner) => {
-			runner.handleBatchStart(batchInfo);
+			runner.handleBatchStart({ ...batchInfo.batchSummary, range: selectionDataToRange(batchInfo.batchSummary.selection) });
 		});
 	}
 
 	public onBatchComplete(batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
 		this._notify(batchInfo.ownerUri, (runner: QueryRunner) => {
-			runner.handleBatchComplete(batchInfo);
+			runner.handleBatchComplete({ range: selectionDataToRange(batchInfo.batchSummary.selection), ...batchInfo.batchSummary });
 		});
 	}
 
 	public onResultSetAvailable(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
 		this._notify(resultSetInfo.ownerUri, (runner: QueryRunner) => {
-			runner.handleResultSetAvailable(resultSetInfo);
+			runner.handleResultSetAvailable(resultSetInfo.resultSetSummary);
 		});
 	}
 
 	public onResultSetUpdated(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
 		this._notify(resultSetInfo.ownerUri, (runner: QueryRunner) => {
-			runner.handleResultSetUpdated(resultSetInfo);
+			runner.handleResultSetUpdated(resultSetInfo.resultSetSummary);
 		});
 	}
 
 	public onMessage(messagesMap: Map<string, azdata.QueryExecuteMessageParams[]>): void {
 		for (const [uri, messages] of messagesMap) {
 			this._notify(uri, (runner: QueryRunner) => {
-				runner.handleMessage(messages);
+				runner.handleMessage(messages.map(m => m.message));
 			});
 		}
 	}
@@ -299,8 +321,8 @@ export class QueryManagementService implements IQueryManagementService {
 	}
 
 	public onEditSessionReady(ownerUri: string, success: boolean, message: string): void {
-		this._notify(ownerUri, (runner: QueryRunner) => {
-			runner.handleEditSessionReady(ownerUri, success, message);
+		this._notify(ownerUri, runner => {
+			(runner as EditQueryRunner).handleEditSessionReady(ownerUri, success, message);
 		});
 	}
 
@@ -352,3 +374,11 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 }
+
+function selectionDataToRange(selection?: azdata.ISelectionData): IRange | undefined {
+	return isUndefined(selection) ? undefined : new Range(selection.startLine + 1, selection.startColumn + 1, selection.endLine + 1, selection.endColumn + 1);
+}
+
+function rangeToSelectionData(range?: IRange): azdata.ISelectionData | undefined {
+	return isUndefined(range) ? undefined : { startLine: range.startLineNumber - 1, startColumn: range.startColumn - 1, endLine: range.endLineNumber - 1, endColumn: range.endColumn - 1 };
+}

src/sql/workbench/services/query/common/queryModel.ts

@@ -3,13 +3,12 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import QueryRunner, { IQueryMessage } from 'sql/workbench/services/query/common/queryRunner';
+import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
+import { IQueryMessage, ResultSetSubset } from 'sql/workbench/services/query/common/query';
 import { DataService } from 'sql/workbench/services/query/common/dataService';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { Event } from 'vs/base/common/event';
 import {
-	ISelectionData,
-	ResultSetSubset,
 	EditUpdateCellResult,
 	EditSessionReadyParams,
 	EditSubsetResult,
@@ -19,6 +18,7 @@ import {
 	queryeditor
 } from 'azdata';
 import { QueryInfo } from 'sql/workbench/services/query/common/queryModelService';
+import { IRange } from 'vs/editor/common/core/range';
 
 export const SERVICE_ID = 'queryModelService';
 
@@ -31,7 +31,7 @@ export interface IQueryPlanInfo {
 }
 
 export interface IQueryInfo {
-	selection: ISelectionData[];
+	range: IRange[];
 	messages: IQueryMessage[];
 }
 
@@ -51,8 +51,8 @@ export interface IQueryModelService {
 	getQueryRunner(uri: string): QueryRunner | undefined;
 
 	getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<ResultSetSubset | undefined>;
-	runQuery(uri: string, selection: ISelectionData | undefined, runOptions?: ExecutionPlanOptions): void;
-	runQueryStatement(uri: string, selection: ISelectionData | undefined): void;
+	runQuery(uri: string, range: IRange | undefined, runOptions?: ExecutionPlanOptions): void;
+	runQueryStatement(uri: string, range: IRange | undefined): void;
 	runQueryString(uri: string, selection: string | undefined): void;
 	cancelQuery(input: QueryRunner | string): void;
 	disposeQuery(uri: string): void;

src/sql/workbench/services/query/common/queryModelService.ts

@@ -5,6 +5,7 @@
 
 import * as GridContentEvents from 'sql/workbench/services/query/common/gridContentEvents';
 import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
+import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
 import { DataService } from 'sql/workbench/services/query/common/dataService';
 import { IQueryModelService, IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
 
@@ -17,6 +18,8 @@ import * as strings from 'vs/base/common/strings';
 import * as types from 'vs/base/common/types';
 import { INotificationService } from 'vs/platform/notification/common/notification';
 import Severity from 'vs/base/common/severity';
+import EditQueryRunner from 'sql/workbench/services/editData/common/editQueryRunner';
+import { IRange } from 'vs/editor/common/core/range';
 
 const selectionSnippetMaxLen = 100;
 
@@ -29,10 +32,10 @@ export interface QueryEvent {
  * Holds information about the state of a query runner
  */
 export class QueryInfo {
-	public queryRunner?: QueryRunner;
+	public queryRunner?: EditQueryRunner;
 	public dataService?: DataService;
 	public queryEventQueue?: QueryEvent[];
-	public selection?: Array<azdata.ISelectionData>;
+	public range?: Array<IRange>;
 	public selectionSnippet?: string;
 
 	// Notes if the angular components have obtained the DataService. If not, all messages sent
@@ -42,7 +45,7 @@ export class QueryInfo {
 	constructor() {
 		this.dataServiceReady = false;
 		this.queryEventQueue = [];
-		this.selection = [];
+		this.range = [];
 	}
 }
 
@@ -128,10 +131,10 @@ export class QueryModelService implements IQueryModelService {
 	/**
 	 * Get more data rows from the current resultSets from the service layer
 	 */
-	public getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<azdata.ResultSetSubset | undefined> {
+	public getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<ResultSetSubset | undefined> {
 		if (this._queryInfoMap.has(uri)) {
 			return this._getQueryInfo(uri)!.queryRunner!.getQueryRows(rowStart, numberOfRows, batchId, resultId).then(results => {
-				return results.resultSubset;
+				return results;
 			});
 		} else {
 			return Promise.resolve(undefined);
@@ -170,15 +173,15 @@ export class QueryModelService implements IQueryModelService {
 	/**
 	 * Run a query for the given URI with the given text selection
 	 */
-	public async runQuery(uri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
-		return this.doRunQuery(uri, selection, false, runOptions);
+	public async runQuery(uri: string, range: IRange, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
+		return this.doRunQuery(uri, range, false, runOptions);
 	}
 
 	/**
 	 * Run the current SQL statement for the given URI
 	 */
-	public async runQueryStatement(uri: string, selection: azdata.ISelectionData): Promise<void> {
-		return this.doRunQuery(uri, selection, true);
+	public async runQueryStatement(uri: string, range: IRange): Promise<void> {
+		return this.doRunQuery(uri, range, true);
 	}
 
 	/**
@@ -191,7 +194,7 @@ export class QueryModelService implements IQueryModelService {
 	/**
 	 * Run Query implementation
 	 */
-	private async doRunQuery(uri: string, selection: azdata.ISelectionData | string,
+	private async doRunQuery(uri: string, range: IRange | string,
 		runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
 		// Reuse existing query runner if it exists
 		let queryRunner: QueryRunner | undefined;
@@ -208,7 +211,7 @@ export class QueryModelService implements IQueryModelService {
 
 			// If the query is not in progress, we can reuse the query runner
 			queryRunner = existingRunner!;
-			info.selection = [];
+			info.range = [];
 			info.selectionSnippet = undefined;
 		} else {
 			// We do not have a query runner for this editor, so create a new one
@@ -217,23 +220,23 @@ export class QueryModelService implements IQueryModelService {
 			queryRunner = info.queryRunner!;
 		}
 
-		if (types.isString(selection)) {
+		if (types.isString(range)) {
 			// Run the query string in this case
-			if (selection.length < selectionSnippetMaxLen) {
-				info.selectionSnippet = selection;
+			if (range.length < selectionSnippetMaxLen) {
+				info.selectionSnippet = range;
 			} else {
-				info.selectionSnippet = selection.substring(0, selectionSnippetMaxLen - 3) + '...';
+				info.selectionSnippet = range.substring(0, selectionSnippetMaxLen - 3) + '...';
 			}
-			return queryRunner.runQuery(selection, runOptions);
+			return queryRunner.runQuery(range, runOptions);
 		} else if (runCurrentStatement) {
-			return queryRunner.runQueryStatement(selection);
+			return queryRunner.runQueryStatement(range);
 		} else {
-			return queryRunner.runQuery(selection, runOptions);
+			return queryRunner.runQuery(range, runOptions);
 		}
 	}
 
 	private initQueryRunner(uri: string): QueryInfo {
-		let queryRunner = this._instantiationService.createInstance(QueryRunner, uri);
+		let queryRunner = this._instantiationService.createInstance(EditQueryRunner, uri);
 		let info = new QueryInfo();
 		queryRunner.onResultSet(e => {
 			this._fireQueryEvent(uri, 'resultSet', e);
@@ -241,14 +244,14 @@ export class QueryModelService implements IQueryModelService {
 		queryRunner.onBatchStart(b => {
 			let link = undefined;
 			let messageText = nls.localize('runQueryBatchStartMessage', "Started executing query at ");
-			if (b.selection) {
+			if (b.range) {
 				if (info.selectionSnippet) {
 					// This indicates it's a query string. Do not include line information since it'll be inaccurate, but show some of the
 					// executed query text
 					messageText = nls.localize('runQueryStringBatchStartMessage', "Started executing query \"{0}\"", info.selectionSnippet);
 				} else {
 					link = {
-						text: strings.format(nls.localize('runQueryBatchStartLine', "Line {0}"), b.selection.startLine + 1)
+						text: strings.format(nls.localize('runQueryBatchStartLine', "Line {0}"), b.range.startLineNumber)
 					};
 				}
 			}
@@ -260,7 +263,7 @@ export class QueryModelService implements IQueryModelService {
 				link: link
 			};
 			this._fireQueryEvent(uri, 'message', message);
-			info.selection!.push(this._validateSelection(b.selection));
+			info.range!.push(b.range);
 		});
 		queryRunner.onMessage(m => {
 			this._fireQueryEvent(uri, 'message', m);
@@ -274,7 +277,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: uri,
 				queryInfo:
 				{
-					selection: info.selection!,
+					range: info.range!,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -292,7 +295,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: uri,
 				queryInfo:
 				{
-					selection: info.selection!,
+					range: info.range!,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -308,7 +311,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: uri,
 				queryInfo:
 				{
-					selection: info.selection!,
+					range: info.range!,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -324,7 +327,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: planInfo.fileUri,
 				queryInfo:
 				{
-					selection: info.selection!,
+					range: info.range!,
 					messages: info.queryRunner!.messages
 				},
 				params: planInfo
@@ -338,7 +341,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: uri,
 				queryInfo:
 				{
-					selection: info.selection!,
+					range: info.range!,
 					messages: info.queryRunner!.messages
 				},
 				params: resultSetInfo
@@ -399,12 +402,12 @@ export class QueryModelService implements IQueryModelService {
 	// EDIT DATA METHODS /////////////////////////////////////////////////////
 	async initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
 		// Reuse existing query runner if it exists
-		let queryRunner: QueryRunner;
+		let queryRunner: EditQueryRunner;
 		let info: QueryInfo;
 
 		if (this._queryInfoMap.has(ownerUri)) {
 			info = this._getQueryInfo(ownerUri)!;
-			let existingRunner: QueryRunner = info.queryRunner!;
+			let existingRunner = info.queryRunner!;
 
 			// If the initialization is already in progress
 			if (existingRunner.isExecuting) {
@@ -417,7 +420,7 @@ export class QueryModelService implements IQueryModelService {
 
 			// We do not have a query runner for this editor, so create a new one
 			// and map it to the results uri
-			queryRunner = this._instantiationService.createInstance(QueryRunner, ownerUri);
+			queryRunner = this._instantiationService.createInstance(EditQueryRunner, ownerUri);
 			const resultSetEventType = 'resultSet';
 			queryRunner.onResultSet(resultSet => {
 				this._fireQueryEvent(ownerUri, resultSetEventType, resultSet);
@@ -428,14 +431,14 @@ export class QueryModelService implements IQueryModelService {
 			queryRunner.onBatchStart(batch => {
 				let link = undefined;
 				let messageText = nls.localize('runQueryBatchStartMessage', "Started executing query at ");
-				if (batch.selection) {
+				if (batch.range) {
 					if (info.selectionSnippet) {
 						// This indicates it's a query string. Do not include line information since it'll be inaccurate, but show some of the
 						// executed query text
 						messageText = nls.localize('runQueryStringBatchStartMessage', "Started executing query \"{0}\"", info.selectionSnippet);
 					} else {
 						link = {
-							text: strings.format(nls.localize('runQueryBatchStartLine', "Line {0}"), batch.selection.startLine + 1)
+							text: strings.format(nls.localize('runQueryBatchStartLine', "Line {0}"), batch.range.startLineNumber)
 						};
 					}
 				}
@@ -459,7 +462,7 @@ export class QueryModelService implements IQueryModelService {
 					uri: ownerUri,
 					queryInfo:
 					{
-						selection: info.selection!,
+						range: info.range!,
 						messages: info.queryRunner!.messages
 					},
 				};
@@ -476,7 +479,7 @@ export class QueryModelService implements IQueryModelService {
 					uri: ownerUri,
 					queryInfo:
 					{
-						selection: info.selection!,
+						range: info.range!,
 						messages: info.queryRunner!.messages
 					},
 				};
@@ -596,8 +599,8 @@ export class QueryModelService implements IQueryModelService {
 
 	// PRIVATE METHODS //////////////////////////////////////////////////////
 
-	private internalGetQueryRunner(ownerUri: string): QueryRunner | undefined {
-		let queryRunner: QueryRunner | undefined;
+	private internalGetQueryRunner(ownerUri: string): EditQueryRunner | undefined {
+		let queryRunner: EditQueryRunner | undefined;
 		if (this._queryInfoMap.has(ownerUri)) {
 			let existingRunner = this._getQueryInfo(ownerUri)!.queryRunner!;
 			// If the query is not already executing then set it up
@@ -648,17 +651,4 @@ export class QueryModelService implements IQueryModelService {
 	public _getQueryInfo(uri: string): QueryInfo | undefined {
 		return this._queryInfoMap.get(uri);
 	}
-
-	// TODO remove this funciton and its usages when #821 in vscode-mssql is fixed and
-	// the SqlToolsService version is updated in this repo - coquagli 4/19/2017
-	private _validateSelection(selection: azdata.ISelectionData): azdata.ISelectionData {
-		if (!selection) {
-			selection = <azdata.ISelectionData>{};
-		}
-		selection.endColumn = selection ? Math.max(0, selection.endColumn) : 0;
-		selection.endLine = selection ? Math.max(0, selection.endLine) : 0;
-		selection.startColumn = selection ? Math.max(0, selection.startColumn) : 0;
-		selection.startLine = selection ? Math.max(0, selection.startLine) : 0;
-		return selection;
-	}
 }

src/sql/workbench/services/query/common/queryRunner.ts

@@ -3,15 +3,12 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import * as azdata from 'azdata';
-
-import { IQueryManagementService } from 'sql/workbench/services/query/common/queryManagement';
+import { IQueryManagementService, QueryCancelResult, ExecutionPlanOptions } from 'sql/workbench/services/query/common/queryManagement';
 import * as Utils from 'sql/platform/connection/common/utils';
 import { Deferred } from 'sql/base/common/promise';
 import { IQueryPlanInfo } from 'sql/workbench/services/query/common/queryModel';
 import { ResultSerializer, SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
 
-import Severity from 'vs/base/common/severity';
 import * as nls from 'vs/nls';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
 import * as types from 'vs/base/common/types';
@@ -27,20 +24,8 @@ import { IGridDataProvider, getResultsString } from 'sql/workbench/services/quer
 import { getErrorMessage } from 'vs/base/common/errors';
 import { ILogService } from 'vs/platform/log/common/log';
 import { find } from 'vs/base/common/arrays';
-
-export interface IEditSessionReadyEvent {
-	ownerUri: string;
-	success: boolean;
-	message: string;
-}
-
-export interface IQueryMessage {
-	batchId?: number;
-	isError: boolean;
-	time?: string;
-	message: string;
-	selection?: azdata.ISelectionData;
-}
+import { IRange, Range } from 'vs/editor/common/core/range';
+import { BatchSummary, IQueryMessage, ResultSetSummary, QueryExecuteSubsetParams, CompleteBatchSummary, IResultMessage, ResultSetSubset, BatchStartSummary } from './query';
 
 /*
 * Query Runner class which handles running a query, reports the results to the content manager,
@@ -50,10 +35,10 @@ export default class QueryRunner extends Disposable {
 	// MEMBER VARIABLES ////////////////////////////////////////////////////
 	private _resultLineOffset?: number;
 	private _resultColumnOffset?: number;
-	private _totalElapsedMilliseconds: number = 0;
-	private _isExecuting: boolean = false;
+	protected _totalElapsedMilliseconds: number = 0;
+	protected _isExecuting: boolean = false;
 	private _hasCompleted: boolean = false;
-	private _batchSets: azdata.BatchSummary[] = [];
+	private _batchSets: BatchSummary[] = [];
 	private _messages: IQueryMessage[] = [];
 	private registered = false;
 
@@ -65,31 +50,28 @@ export default class QueryRunner extends Disposable {
 	private _onMessage = this._register(new Emitter<IQueryMessage[]>());
 	public get onMessage(): Event<IQueryMessage[]> { return this._onMessage.event; } // this is the only way typemoq can moq this... needs investigation @todo anthonydresser 5/2/2019
 
-	private _onResultSet = this._register(new Emitter<azdata.ResultSetSummary>());
+	private readonly _onResultSet = this._register(new Emitter<ResultSetSummary>());
 	public readonly onResultSet = this._onResultSet.event;
 
-	private _onResultSetUpdate = this._register(new Emitter<azdata.ResultSetSummary>());
+	private readonly _onResultSetUpdate = this._register(new Emitter<ResultSetSummary>());
 	public readonly onResultSetUpdate = this._onResultSetUpdate.event;
 
-	private _onQueryStart = this._register(new Emitter<void>());
+	protected readonly _onQueryStart = this._register(new Emitter<void>());
 	public readonly onQueryStart: Event<void> = this._onQueryStart.event;
 
-	private _onQueryEnd = this._register(new Emitter<string>());
+	private readonly _onQueryEnd = this._register(new Emitter<string>());
 	public get onQueryEnd(): Event<string> { return this._onQueryEnd.event; }
 
-	private _onBatchStart = this._register(new Emitter<azdata.BatchSummary>());
-	public readonly onBatchStart: Event<azdata.BatchSummary> = this._onBatchStart.event;
+	private readonly _onBatchStart = this._register(new Emitter<BatchStartSummary>());
+	public readonly onBatchStart: Event<BatchStartSummary> = this._onBatchStart.event;
 
-	private _onBatchEnd = this._register(new Emitter<azdata.BatchSummary>());
-	public readonly onBatchEnd: Event<azdata.BatchSummary> = this._onBatchEnd.event;
+	private readonly _onBatchEnd = this._register(new Emitter<CompleteBatchSummary>());
+	public readonly onBatchEnd: Event<CompleteBatchSummary> = this._onBatchEnd.event;
 
-	private _onEditSessionReady = this._register(new Emitter<IEditSessionReadyEvent>());
-	public readonly onEditSessionReady = this._onEditSessionReady.event;
-
-	private _onQueryPlanAvailable = this._register(new Emitter<IQueryPlanInfo>());
+	private readonly _onQueryPlanAvailable = this._register(new Emitter<IQueryPlanInfo>());
 	public readonly onQueryPlanAvailable = this._onQueryPlanAvailable.event;
 
-	private _onVisualize = this._register(new Emitter<azdata.ResultSetSummary>());
+	private readonly _onVisualize = this._register(new Emitter<ResultSetSummary>());
 	public readonly onVisualize = this._onVisualize.event;
 
 	private _queryStartTime?: Date;
@@ -104,12 +86,11 @@ export default class QueryRunner extends Disposable {
 	// CONSTRUCTOR /////////////////////////////////////////////////////////
 	constructor(
 		public uri: string,
-		@IQueryManagementService private _queryManagementService: IQueryManagementService,
-		@INotificationService private _notificationService: INotificationService,
-		@IConfigurationService private _configurationService: IConfigurationService,
-		@IInstantiationService private instantiationService: IInstantiationService,
-		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService,
-		@ILogService private _logService: ILogService
+		@IQueryManagementService protected readonly queryManagementService: IQueryManagementService,
+		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@IInstantiationService private readonly instantiationService: IInstantiationService,
+		@ITextResourcePropertiesService private readonly textResourcePropertiesService: ITextResourcePropertiesService,
+		@ILogService private readonly logService: ILogService
 	) {
 		super();
 	}
@@ -125,7 +106,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * For public use only, for private use, directly access the member
 	 */
-	public get batchSets(): azdata.BatchSummary[] {
+	public get batchSets(): BatchSummary[] {
 		return this._batchSets.slice(0);
 	}
 
@@ -141,22 +122,22 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Cancels the running query, if there is one
 	 */
-	public cancelQuery(): Promise<azdata.QueryCancelResult> {
-		return this._queryManagementService.cancelQuery(this.uri);
+	public cancelQuery(): Promise<QueryCancelResult> {
+		return this.queryManagementService.cancelQuery(this.uri);
 	}
 
 	/**
 	 * Runs the query with the provided query
 	 * @param input Query string to execute
 	 */
-	public runQuery(input: string, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	public runQuery(input: string, runOptions?: ExecutionPlanOptions): Promise<void>;
 	/**
 	 * Runs the query by pulling the query from the document using the provided selection data
 	 * @param input selection data
 	 */
-	public runQuery(input: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
-	public runQuery(input: string | azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
-		if (types.isString(input)) {
+	public runQuery(input: IRange | undefined, runOptions?: ExecutionPlanOptions): Promise<void>;
+	public runQuery(input: string | IRange | undefined, runOptions?: ExecutionPlanOptions): Promise<void> {
+		if (types.isString(input) || types.isUndefined(input)) {
 			return this.doRunQuery(input, false, runOptions);
 		} else {
 			return this.doRunQuery(input, false, runOptions);
@@ -167,7 +148,7 @@ export default class QueryRunner extends Disposable {
 	 * Runs the current SQL statement by pulling the query from the document using the provided selection data
 	 * @param input selection data
 	 */
-	public runQueryStatement(input: azdata.ISelectionData): Promise<void> {
+	public runQueryStatement(input: IRange): Promise<void> {
 		return this.doRunQuery(input, true);
 	}
 
@@ -175,9 +156,9 @@ export default class QueryRunner extends Disposable {
 	 * Implementation that runs the query with the provided query
 	 * @param input Query string to execute
 	 */
-	private doRunQuery(input: string, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
-	private doRunQuery(input: azdata.ISelectionData, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
-	private doRunQuery(input: string | azdata.ISelectionData, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
+	private doRunQuery(input: string, runCurrentStatement: boolean, runOptions?: ExecutionPlanOptions): Promise<void>;
+	private doRunQuery(input: IRange | undefined, runCurrentStatement: boolean, runOptions?: ExecutionPlanOptions): Promise<void>;
+	private doRunQuery(input: string | IRange | undefined, runCurrentStatement: boolean, runOptions?: ExecutionPlanOptions): Promise<void> {
 		if (this.isExecuting) {
 			return Promise.resolve();
 		}
@@ -187,9 +168,9 @@ export default class QueryRunner extends Disposable {
 		this._queryStartTime = undefined;
 		this._queryEndTime = undefined;
 		this._messages = [];
-		if (isSelectionOrUndefined(input)) {
+		if (isRangeOrUndefined(input)) {
 			// Update internal state to show that we're executing the query
-			this._resultLineOffset = input ? input.startLine : 0;
+			this._resultLineOffset = input ? input.startLineNumber : 0;
 			this._resultColumnOffset = input ? input.startColumn : 0;
 			this._isExecuting = true;
 			this._totalElapsedMilliseconds = 0;
@@ -199,8 +180,8 @@ export default class QueryRunner extends Disposable {
 
 			// Send the request to execute the query
 			return runCurrentStatement
-				? this._queryManagementService.runQueryStatement(this.uri, input.startLine, input.startColumn).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e))
-				: this._queryManagementService.runQuery(this.uri, input, runOptions).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
+				? this.queryManagementService.runQueryStatement(this.uri, input.startLineNumber, input.startColumn).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e))
+				: this.queryManagementService.runQuery(this.uri, input, runOptions).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
 		} else {
 			// Update internal state to show that we're executing the query
 			this._isExecuting = true;
@@ -208,7 +189,7 @@ export default class QueryRunner extends Disposable {
 
 			this._onQueryStart.fire();
 
-			return this._queryManagementService.runQueryString(this.uri, input).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
+			return this.queryManagementService.runQueryString(this.uri, input).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
 		}
 	}
 
@@ -218,45 +199,39 @@ export default class QueryRunner extends Disposable {
 		// The query has started, so lets fire up the result pane
 		if (!this.registered) {
 			this.registered = true;
-			this._queryManagementService.registerRunner(this, this.uri);
+			this.queryManagementService.registerRunner(this, this.uri);
 		}
 	}
 
 	private handleFailureRunQueryResult(error: any) {
 		// Attempting to launch the query failed, show the error message
-		const eol = getEolString(this._textResourcePropertiesService, this.uri);
+		const eol = getEolString(this.textResourcePropertiesService, this.uri);
 		if (error instanceof Error) {
 			error = error.message;
 		}
 		let message = nls.localize('query.ExecutionFailedError', "Execution failed due to an unexpected error: {0}\t{1}", eol, error);
-		this.handleMessage([<azdata.QueryExecuteMessageParams>{
-			ownerUri: this.uri,
-			message: {
-				isError: true,
-				message: message
-			}
+		this.handleMessage([{
+			isError: true,
+			message: message
 		}]);
-		this.handleQueryComplete(<azdata.QueryExecuteCompleteNotificationResult>{ ownerUri: this.uri });
+		this.handleQueryComplete();
 	}
 
 	/**
 	 * Handle a QueryComplete from the service layer
 	 */
-	public handleQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void {
+	public handleQueryComplete(batchSummaries?: CompleteBatchSummary[]): void {
 		// this also isn't exact but its the best we can do
 		this._queryEndTime = new Date();
 
 		// Store the batch sets we got back as a source of "truth"
 		this._isExecuting = false;
 		this._hasCompleted = true;
-		this._batchSets = result.batchSummaries ? result.batchSummaries : [];
+		this._batchSets = batchSummaries ? batchSummaries : [];
 
 		this._batchSets.map(batch => {
-			if (batch.selection) {
-				batch.selection.startLine += this._resultLineOffset!;
-				batch.selection.startColumn += this._resultColumnOffset!;
-				batch.selection.endLine += this._resultLineOffset!;
-				batch.selection.endColumn += this._resultColumnOffset!;
+			if (batch.range) {
+				batch.range = new Range(batch.range.startLineNumber + this._resultLineOffset, batch.range.startColumn + this._resultColumnOffset, batch.range.endLineNumber + this._resultLineOffset, batch.range.endColumn + this._resultColumnOffset);
 			}
 		});
 
@@ -277,28 +252,20 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Handle a BatchStart from the service layer
 	 */
-	public handleBatchStart(result: azdata.QueryExecuteBatchNotificationParams): void {
-		let batch = result.batchSummary;
-
+	public handleBatchStart(batch: BatchStartSummary): void {
 		// Recalculate the start and end lines, relative to the result line offset
-		if (batch.selection) {
-			batch.selection.startLine += this._resultLineOffset!;
-			batch.selection.startColumn += this._resultColumnOffset!;
-			batch.selection.endLine += this._resultLineOffset!;
-			batch.selection.endColumn += this._resultColumnOffset!;
+		if (batch.range) {
+			batch.range = new Range(batch.range.startLineNumber + this._resultLineOffset, batch.range.startColumn + this._resultColumnOffset, batch.range.endLineNumber + this._resultLineOffset, batch.range.endColumn + this._resultColumnOffset);
 		}
 
-		// Set the result sets as an empty array so that as result sets complete we can add to the list
-		batch.resultSetSummaries = [];
-
 		// Store the batch
-		this._batchSets[batch.id] = batch;
+		this._batchSets[batch.id] = { ...batch, resultSetSummaries: [], hasError: false };
 
-		let message = {
+		let message: IQueryMessage = {
 			// account for index by 1
-			message: nls.localize('query.message.startQuery', "Started executing query at Line {0}", batch.selection.startLine + 1),
+			message: batch.range ? nls.localize('query.message.startQueryWithRange', "Started executing query at Line {0}", batch.range.startLineNumber) : nls.localize('query.message.startQuery', "Started executing batch {0}", batch.id),
 			time: batch.executionStart,
-			selection: batch.selection,
+			range: batch.range,
 			isError: false
 		};
 		this._messages.push(message);
@@ -309,9 +276,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Handle a BatchComplete from the service layer
 	 */
-	public handleBatchComplete(result: azdata.QueryExecuteBatchNotificationParams): void {
-		let batch: azdata.BatchSummary = result.batchSummary;
-
+	public handleBatchComplete(batch: CompleteBatchSummary): void {
 		// Store the batch again to get the rest of the data
 		this._batchSets[batch.id] = batch;
 		let executionTime = <number>(Utils.parseTimeString(batch.executionElapsed) || 0);
@@ -327,19 +292,18 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Handle a ResultSetComplete from the service layer
 	 */
-	public handleResultSetAvailable(result: azdata.QueryExecuteResultSetNotificationParams): void {
-		if (result && result.resultSetSummary) {
-			let resultSet = result.resultSetSummary;
-			let batchSet: azdata.BatchSummary;
+	public handleResultSetAvailable(resultSet?: ResultSetSummary): void {
+		if (resultSet) {
+			let batchSet: BatchSummary;
 			if (!resultSet.batchId) {
 				// Missing the batchId or processing batchId==0. In this case, default to always using the first batch in the list
 				// or create one in the case the DMP extension didn't obey the contract perfectly
 				if (this._batchSets.length > 0) {
 					batchSet = this._batchSets[0];
 				} else {
-					batchSet = <azdata.BatchSummary><unknown>{
+					batchSet = <BatchSummary>{
 						id: 0,
-						selection: undefined,
+						range: undefined,
 						hasError: false,
 						resultSetSummaries: []
 					};
@@ -350,15 +314,15 @@ export default class QueryRunner extends Disposable {
 			}
 			// handle getting queryPlanxml if we need too
 			// check if this result has show plan, this needs work, it won't work for any other provider
-			let hasShowPlan = !!find(result.resultSetSummary.columnInfo, e => e.columnName === 'Microsoft SQL Server 2005 XML Showplan');
-			if (hasShowPlan) {
+			let hasShowPlan = !!find(resultSet.columnInfo, e => e.columnName === 'Microsoft SQL Server 2005 XML Showplan');
+			if (hasShowPlan && resultSet.rowCount > 0) {
 				this._isQueryPlan = true;
 
-				this.getQueryRows(0, 1, result.resultSetSummary.batchId, result.resultSetSummary.id).then(e => {
-					if (e.resultSubset.rows) {
-						this._planXml.resolve(e.resultSubset.rows[0][0].displayValue);
+				this.getQueryRows(0, 1, resultSet.batchId, resultSet.id).then(e => {
+					if (e.rows) {
+						this._planXml.resolve(e.rows[0][0].displayValue);
 					}
-				}).catch((e) => this._logService.error(e));
+				}).catch((e) => this.logService.error(e));
 			}
 			// we will just ignore the set if we already have it
 			// ideally this should never happen
@@ -370,31 +334,30 @@ export default class QueryRunner extends Disposable {
 		}
 	}
 
-	public handleResultSetUpdated(result: azdata.QueryExecuteResultSetNotificationParams): void {
-		if (result && result.resultSetSummary) {
-			let resultSet = result.resultSetSummary;
-			let batchSet: azdata.BatchSummary;
+	public handleResultSetUpdated(resultSet?: ResultSetSummary): void {
+		if (resultSet) {
+			let batchSet: BatchSummary;
 			batchSet = this._batchSets[resultSet.batchId];
 			// handle getting queryPlanxml if we need too
 			// check if this result has show plan, this needs work, it won't work for any other provider
-			let hasShowPlan = !!find(result.resultSetSummary.columnInfo, e => e.columnName === 'Microsoft SQL Server 2005 XML Showplan');
+			let hasShowPlan = !!resultSet.columnInfo.find(e => e.columnName === 'Microsoft SQL Server 2005 XML Showplan');
 			if (hasShowPlan) {
 				this._isQueryPlan = true;
-				this.getQueryRows(0, 1, result.resultSetSummary.batchId, result.resultSetSummary.id).then(e => {
+				this.getQueryRows(0, 1, resultSet.batchId, resultSet.id).then(e => {
 
-					if (e.resultSubset.rows) {
-						let planXmlString = e.resultSubset.rows[0][0].displayValue;
-						this._planXml.resolve(e.resultSubset.rows[0][0].displayValue);
+					if (e.rows) {
+						let planXmlString = e.rows[0][0].displayValue;
+						this._planXml.resolve(e.rows[0][0].displayValue);
 						// fire query plan available event if execution is completed
-						if (result.resultSetSummary.complete) {
+						if (resultSet.complete) {
 							this._onQueryPlanAvailable.fire({
 								providerId: mssqlProviderName,
-								fileUri: result.ownerUri,
+								fileUri: this.uri,
 								planXml: planXmlString
 							});
 						}
 					}
-				}).catch((e) => this._logService.error(e));
+				}).catch((e) => this.logService.error(e));
 			}
 			if (batchSet) {
 				// Store the result set in the batch and emit that a result set has completed
@@ -407,8 +370,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Handle a Mssage from the service layer
 	 */
-	public handleMessage(messagesObj: azdata.QueryExecuteMessageParams[]): void {
-		const messages = messagesObj.map(m => m.message);
+	public handleMessage(messages: IResultMessage[]): void {
 		this._messages.push(...messages);
 
 		// Send the message to the results pane
@@ -418,8 +380,8 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Get more data rows from the current resultSets from the service layer
 	 */
-	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Promise<azdata.QueryExecuteSubsetResult> {
-		let rowData: azdata.QueryExecuteSubsetParams = <azdata.QueryExecuteSubsetParams>{
+	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Promise<ResultSetSubset> {
+		let rowData: QueryExecuteSubsetParams = <QueryExecuteSubsetParams>{
 			ownerUri: this.uri,
 			resultSetIndex: resultSetIndex,
 			rowsCount: numberOfRows,
@@ -427,7 +389,7 @@ export default class QueryRunner extends Disposable {
 			batchIndex: batchIndex
 		};
 
-		return this._queryManagementService.getQueryRows(rowData).then(r => r, error => {
+		return this.queryManagementService.getQueryRows(rowData).then(r => r, error => {
 			// this._notificationService.notify({
 			// 	severity: Severity.Error,
 			// 	message: nls.localize('query.gettingRowsFailedError', 'Something went wrong getting more rows: {0}', error)
@@ -436,104 +398,11 @@ export default class QueryRunner extends Disposable {
 		});
 	}
 
-	/*
-	 * Handle a session ready event for Edit Data
-	 */
-	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
-		// Update internal state to show that we're executing the query
-		this._isExecuting = true;
-		this._totalElapsedMilliseconds = 0;
-		// TODO issue #228 add statusview callbacks here
-
-		return this._queryManagementService.initializeEdit(ownerUri, schemaName, objectName, objectType, rowLimit, queryString).then(result => {
-			// The query has started, so lets fire up the result pane
-			this._onQueryStart.fire();
-			this._queryManagementService.registerRunner(this, ownerUri);
-		}, error => {
-			// Attempting to launch the query failed, show the error message
-
-			// TODO issue #228 add statusview callbacks here
-			this._isExecuting = false;
-			this._notificationService.error(nls.localize('query.initEditExecutionFailed', "Initialize edit data session failed: ") + error);
-		});
-	}
-
-	/**
-	 * Retrieves a number of rows from an edit session
-	 * @param rowStart     The index of the row to start returning (inclusive)
-	 * @param numberOfRows The number of rows to return
-	 */
-	public getEditRows(rowStart: number, numberOfRows: number): Promise<azdata.EditSubsetResult> {
-		const self = this;
-		let rowData: azdata.EditSubsetParams = {
-			ownerUri: this.uri,
-			rowCount: numberOfRows,
-			rowStartIndex: rowStart
-		};
-
-		return new Promise<azdata.EditSubsetResult>((resolve, reject) => {
-			self._queryManagementService.getEditRows(rowData).then(result => {
-				if (!result.hasOwnProperty('rowCount')) {
-					let error = `Nothing returned from subset query`;
-					self._notificationService.notify({
-						severity: Severity.Error,
-						message: error
-					});
-					reject(error);
-				}
-				resolve(result);
-			}, error => {
-				// let errorMessage = nls.localize('query.moreRowsFailedError', "Something went wrong getting more rows:");
-				// self._notificationService.notify({
-				// 	severity: Severity.Error,
-				// 	message: `${errorMessage} ${error}`
-				// });
-				reject(error);
-			});
-		});
-	}
-
-	public handleEditSessionReady(ownerUri: string, success: boolean, message: string): void {
-		this._onEditSessionReady.fire({ ownerUri, success, message });
-	}
-
-	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult> {
-		return this._queryManagementService.updateCell(ownerUri, rowId, columnId, newValue);
-	}
-
-	public commitEdit(ownerUri: string): Promise<void> {
-		return this._queryManagementService.commitEdit(ownerUri);
-	}
-
-	public createRow(ownerUri: string): Promise<azdata.EditCreateRowResult> {
-		return this._queryManagementService.createRow(ownerUri).then(result => {
-			return result;
-		});
-	}
-
-	public deleteRow(ownerUri: string, rowId: number): Promise<void> {
-		return this._queryManagementService.deleteRow(ownerUri, rowId);
-	}
-
-	public revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult> {
-		return this._queryManagementService.revertCell(ownerUri, rowId, columnId).then(result => {
-			return result;
-		});
-	}
-
-	public revertRow(ownerUri: string, rowId: number): Promise<void> {
-		return this._queryManagementService.revertRow(ownerUri, rowId);
-	}
-
-	public disposeEdit(ownerUri: string): Promise<void> {
-		return this._queryManagementService.disposeEdit(ownerUri);
-	}
-
 	/**
 	 * Disposes the Query from the service client
 	 */
 	public async disposeQuery(): Promise<void> {
-		await this._queryManagementService.disposeQuery(this.uri);
+		await this.queryManagementService.disposeQuery(this.uri);
 		this.dispose();
 	}
 
@@ -561,7 +430,7 @@ export default class QueryRunner extends Disposable {
 
 	public getColumnHeaders(batchId: number, resultId: number, range: Slick.Range): string[] | undefined {
 		let headers: string[] | undefined = undefined;
-		let batchSummary: azdata.BatchSummary = this._batchSets[batchId];
+		let batchSummary: BatchSummary = this._batchSets[batchId];
 		if (batchSummary !== undefined) {
 			let resultSetSummary = batchSummary.resultSetSummaries[resultId];
 			headers = resultSetSummary.columnInfo.slice(range.fromCell, range.toCell + 1).map((info, i) => {
@@ -573,7 +442,7 @@ export default class QueryRunner extends Disposable {
 
 	private sendBatchTimeMessage(batchId: number, executionTime: string): void {
 		// get config copyRemoveNewLine option from vscode config
-		let showBatchTime = this._configurationService.getValue<boolean>('sql.showBatchTime');
+		let showBatchTime = this.configurationService.getValue<boolean>('sql.showBatchTime');
 		if (showBatchTime) {
 			let message: IQueryMessage = {
 				batchId: batchId,
@@ -596,7 +465,7 @@ export default class QueryRunner extends Disposable {
 	}
 
 	public notifyVisualizeRequested(batchId: number, resultSetId: number): void {
-		let result: azdata.ResultSetSummary = {
+		let result: ResultSetSummary = {
 			batchId: batchId,
 			id: resultSetId,
 			columnInfo: this.batchSets[batchId].resultSetSummaries[resultSetId].columnInfo,
@@ -620,7 +489,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	) {
 	}
 
-	getRowData(rowStart: number, numberOfRows: number): Promise<azdata.QueryExecuteSubsetResult> {
+	getRowData(rowStart: number, numberOfRows: number): Promise<ResultSetSubset> {
 		return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId);
 	}
 
@@ -679,6 +548,6 @@ export function shouldRemoveNewLines(configurationService: IConfigurationService
 	return !!removeNewLines;
 }
 
-function isSelectionOrUndefined(input: string | azdata.ISelectionData | undefined): input is azdata.ISelectionData | undefined {
-	return types.isObject(input) || types.isUndefinedOrNull(input);
+function isRangeOrUndefined(input: string | IRange | undefined): input is IRange | undefined {
+	return Range.isIRange(input) || types.isUndefinedOrNull(input);
 }