Revert "Remove references to query model service in query code (#6513)" (#6591)

This reverts commit 674351dc75.

src/sql/workbench/common/customInputConverter.ts

@@ -43,7 +43,9 @@ export function convertEditorInput(input: EditorInput, options: IQueryEditorOpti
 			//QueryInput
 			uri = getQueryEditorFileUri(input);
 			if (uri) {
-				return instantiationService.createInstance(QueryInput, '', input, undefined);
+				const queryResultsInput: QueryResultsInput = instantiationService.createInstance(QueryResultsInput, uri.toString());
+				let queryInput: QueryInput = instantiationService.createInstance(QueryInput, '', input, queryResultsInput, undefined);
+				return queryInput;
 			}
 
 			//QueryPlanInput

src/sql/workbench/parts/commandLine/test/electron-browser/commandLine.test.ts

@@ -26,9 +26,6 @@ import { QueryInput, QueryEditorState } from 'sql/workbench/parts/query/common/q
 import { URI } from 'vs/base/common/uri';
 import { ILogService, NullLogService } from 'vs/platform/log/common/log';
 import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';
-import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
-import { TestEditorInput } from 'vs/workbench/services/editor/test/browser/editorGroupsService.test';
-import { Schemas } from 'vs/base/common/network';
 
 class TestParsedArgs implements ParsedArgs {
 	[arg: string]: any;
@@ -375,9 +372,7 @@ suite('commandLineService tests', () => {
 			}).verifiable(TypeMoq.Times.once());
 		connectionManagementService.setup(c => c.getConnectionProfileById(TypeMoq.It.isAnyString())).returns(() => originalProfile);
 		const configurationService = getConfigurationServiceMock(true);
-		const instantiationService = new TestInstantiationService();
+		const queryInput: TypeMoq.Mock<QueryInput> = TypeMoq.Mock.ofType<QueryInput>(QueryInput);
-		const fileInput = new TestEditorInput(URI.from({ scheme: Schemas.untitled }));
-		const queryInput: TypeMoq.Mock<QueryInput> = TypeMoq.Mock.ofType<QueryInput>(QueryInput, TypeMoq.MockBehavior.Loose, undefined, fileInput, undefined, undefined, undefined, instantiationService);
 		let uri = URI.file(args._[0]);
 		const queryState = new QueryEditorState();
 		queryState.connected = true;

src/sql/workbench/parts/query/browser/keyboardQueryActions.ts

@@ -13,6 +13,7 @@ import * as azdata from 'azdata';
 import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { QueryEditor } from 'sql/workbench/parts/query/browser/queryEditor';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
 import * as WorkbenchUtils from 'sql/workbench/common/sqlWorkbenchUtils';
 import * as Constants from 'sql/workbench/parts/query/common/constants';
 import * as ConnectionConstants from 'sql/platform/connection/common/constants';
@@ -238,6 +239,7 @@ export class RunQueryShortcutAction extends Action {
 
 	constructor(
 		@IEditorService private readonly editorService: IEditorService,
+		@IQueryModelService protected readonly queryModelService: IQueryModelService,
 		@IQueryManagementService private readonly queryManagementService: IQueryManagementService,
 		@IConnectionManagementService private readonly connectionManagementService: IConnectionManagementService,
 		@IConfigurationService private readonly configurationService: IConfigurationService

src/sql/workbench/parts/query/browser/queryActions.ts

@@ -24,6 +24,7 @@ import {
 	RunQueryOnConnectionMode
 } from 'sql/platform/connection/common/connectionManagement';
 import { QueryEditor } from 'sql/workbench/parts/query/browser/queryEditor';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
 import { SelectBox } from 'sql/base/browser/ui/selectBox/selectBox';
 import { attachEditableDropdownStyler, attachSelectBoxStyler } from 'sql/platform/theme/common/styler';
 import { Dropdown } from 'sql/base/parts/editableDropdown/browser/dropdown';
@@ -106,7 +107,9 @@ export class RunQueryAction extends QueryTaskbarAction {
 
 	constructor(
 		editor: QueryEditor,
-		@IConnectionManagementService connectionManagementService: IConnectionManagementService
+		@IQueryModelService protected readonly queryModelService: IQueryModelService,
+		@IConnectionManagementService connectionManagementService: IConnectionManagementService,
+		@IExtensionTipsService private readonly extensionTipsService: IExtensionTipsService
 	) {
 		super(connectionManagementService, editor, RunQueryAction.ID, RunQueryAction.EnabledClass);
 		this.label = nls.localize('runQueryLabel', "Run");
@@ -175,7 +178,8 @@ export class CancelQueryAction extends QueryTaskbarAction {
 
 	constructor(
 		editor: QueryEditor,
-		@IConnectionManagementService connectionManagementService: IConnectionManagementService,
+		@IQueryModelService private readonly queryModelService: IQueryModelService,
+		@IConnectionManagementService connectionManagementService: IConnectionManagementService
 	) {
 		super(connectionManagementService, editor, CancelQueryAction.ID, CancelQueryAction.EnabledClass);
 		this.enabled = false;
@@ -184,7 +188,7 @@ export class CancelQueryAction extends QueryTaskbarAction {
 
 	public run(): Promise<void> {
 		if (this.isConnected(this.editor)) {
-			this.editor.input.runner.cancelQuery();
+			this.queryModelService.cancelQuery(this.editor.input.uri);
 		}
 		return Promise.resolve(null);
 	}

src/sql/workbench/parts/query/browser/queryResultsView.ts

@@ -5,6 +5,7 @@
 
 import { QueryResultsInput } from 'sql/workbench/parts/query/common/queryResultsInput';
 import { TabbedPanel, IPanelTab, IPanelView } from 'sql/base/browser/ui/panel/panel';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
 import QueryRunner from 'sql/platform/query/common/queryRunner';
 import { MessagePanel } from 'sql/workbench/parts/query/browser/messagePanel';
 import { GridPanel } from 'sql/workbench/parts/query/browser/gridPanel';
@@ -178,6 +179,7 @@ export class QueryResultsView extends Disposable {
 		container: HTMLElement,
 		@IThemeService themeService: IThemeService,
 		@IInstantiationService private instantiationService: IInstantiationService,
+		@IQueryModelService private queryModelService: IQueryModelService
 	) {
 		super();
 		this.resultsTab = this._register(new ResultsTab(instantiationService));
@@ -307,7 +309,19 @@ export class QueryResultsView extends Disposable {
 			dynamicTab.captureState(this.input.state.dynamicModelViewTabsState);
 		});
 
-		this.setQueryRunner(input.runner);
+		let info = this.queryModelService._getQueryInfo(input.uri);
+		if (info) {
+			this.setQueryRunner(info.queryRunner);
+		} else {
+			let disposable = this.queryModelService.onRunQueryStart(c => {
+				if (c === input.uri) {
+					let info = this.queryModelService._getQueryInfo(input.uri);
+					this.setQueryRunner(info.queryRunner);
+					disposable.dispose();
+				}
+			});
+			this.runnerDisposables.push(disposable);
+		}
 	}
 
 	clearInput() {

src/sql/workbench/parts/query/browser/statusBarItems.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
 import { IntervalTimer } from 'vs/base/common/async';
 import { IStatusbarService, StatusbarAlignment, IStatusbarEntryAccessor } from 'vs/platform/statusbar/common/statusbar';
 import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
@@ -26,6 +27,7 @@ export class TimeElapsedStatusBarContributions extends Disposable implements IWo
 	constructor(
 		@IStatusbarService private readonly statusbarService: IStatusbarService,
 		@IEditorService private readonly editorService: IEditorService,
+		@IQueryModelService private readonly queryModelService: IQueryModelService
 	) {
 		super();
 		this.statusItem = this._register(
@@ -55,16 +57,30 @@ export class TimeElapsedStatusBarContributions extends Disposable implements IWo
 		this.hide();
 		const activeInput = this.editorService.activeEditor;
 		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
-			const runner = activeInput.runner;
+			const uri = activeInput.uri;
-			if (runner.hasCompleted || runner.isExecuting) {
+			const runner = this.queryModelService.getQueryRunner(uri);
-				this._displayValue(runner);
+			if (runner) {
+				if (runner.hasCompleted || runner.isExecuting) {
+					this._displayValue(runner);
+				}
+				this.disposable.add(runner.onQueryStart(e => {
+					this._displayValue(runner);
+				}));
+				this.disposable.add(runner.onQueryEnd(e => {
+					this._displayValue(runner);
+				}));
+			} else {
+				this.disposable.add(this.queryModelService.onRunQueryStart(e => {
+					if (e === uri) {
+						this._displayValue(this.queryModelService.getQueryRunner(uri));
+					}
+				}));
+				this.disposable.add(this.queryModelService.onRunQueryComplete(e => {
+					if (e === uri) {
+						this._displayValue(this.queryModelService.getQueryRunner(uri));
+					}
+				}));
 			}
-			this.disposable.add(runner.onQueryStart(e => {
-				this._displayValue(runner);
-			}));
-			this.disposable.add(runner.onQueryEnd(e => {
-				this._displayValue(runner);
-			}));
 		}
 	}
 
@@ -103,7 +119,8 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 
 	constructor(
 		@IStatusbarService private readonly statusbarService: IStatusbarService,
-		@IEditorService private readonly editorService: IEditorService
+		@IEditorService private readonly editorService: IEditorService,
+		@IQueryModelService private readonly queryModelService: IQueryModelService
 	) {
 		super();
 		this.statusItem = this._register(
@@ -133,16 +150,29 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 		const activeInput = this.editorService.activeEditor;
 		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
 			const uri = activeInput.uri;
-			const runner = activeInput.runner;
+			const runner = this.queryModelService.getQueryRunner(uri);
-			if (runner.hasCompleted || runner.isExecuting) {
+			if (runner) {
-				this._displayValue(runner);
+				if (runner.hasCompleted || runner.isExecuting) {
+					this._displayValue(runner);
+				}
+				this.disposable.add(runner.onQueryStart(e => {
+					this._displayValue(runner);
+				}));
+				this.disposable.add(runner.onQueryEnd(e => {
+					this._displayValue(runner);
+				}));
+			} else {
+				this.disposable.add(this.queryModelService.onRunQueryStart(e => {
+					if (e === uri) {
+						this._displayValue(this.queryModelService.getQueryRunner(uri));
+					}
+				}));
+				this.disposable.add(this.queryModelService.onRunQueryComplete(e => {
+					if (e === uri) {
+						this._displayValue(this.queryModelService.getQueryRunner(uri));
+					}
+				}));
 			}
-			this.disposable.add(runner.onQueryStart(e => {
-				this._displayValue(runner);
-			}));
-			this.disposable.add(runner.onQueryEnd(e => {
-				this._displayValue(runner);
-			}));
 		}
 	}
 
@@ -161,11 +191,13 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 export class QueryStatusStatusBarContributions extends Disposable implements IWorkbenchContribution {
 
 	private static readonly ID = 'status.query.status';
-	private runnerDisposables = new DisposableStore();
+
+	private visisbleUri: string | undefined;
 
 	constructor(
 		@IStatusbarService private readonly statusbarService: IStatusbarService,
-		@IEditorService private readonly editorService: IEditorService
+		@IEditorService private readonly editorService: IEditorService,
+		@IQueryModelService private readonly queryModelService: IQueryModelService
 	) {
 		super();
 		this._register(
@@ -177,21 +209,22 @@ export class QueryStatusStatusBarContributions extends Disposable implements IWo
 				StatusbarAlignment.RIGHT, 100)
 		);
 
+		this._register(Event.filter(this.queryModelService.onRunQueryStart, uri => uri === this.visisbleUri)(this.update, this));
+		this._register(Event.filter(this.queryModelService.onRunQueryComplete, uri => uri === this.visisbleUri)(this.update, this));
 		this._register(this.editorService.onDidActiveEditorChange(this.update, this));
 		this.update();
 	}
 
 	private update() {
 		this.hide();
-		this.runnerDisposables.clear();
+		this.visisbleUri = undefined;
 		const activeInput = this.editorService.activeEditor;
 		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
-			const runner = activeInput.runner;
+			this.visisbleUri = activeInput.uri;
+			const runner = this.queryModelService.getQueryRunner(this.visisbleUri);
 			if (runner && runner.isExecuting) {
 				this.show();
 			}
-			this.runnerDisposables.add(runner.onQueryStart(() => this.show()));
-			this.runnerDisposables.add(runner.onQueryEnd(() => this.hide()));
 		}
 	}
 

src/sql/workbench/parts/query/common/queryInput.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { localize } from 'vs/nls';
-import { IDisposable, Disposable } from 'vs/base/common/lifecycle';
+import { IDisposable, dispose, Disposable } from 'vs/base/common/lifecycle';
 import { Event, Emitter } from 'vs/base/common/event';
 import { URI } from 'vs/base/common/uri';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
@@ -13,13 +13,11 @@ import { IConfigurationService } from 'vs/platform/configuration/common/configur
 
 import { IConnectionManagementService, IConnectableInput, INewConnectionParams, RunQueryOnConnectionMode } from 'sql/platform/connection/common/connectionManagement';
 import { QueryResultsInput } from 'sql/workbench/parts/query/common/queryResultsInput';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
 
 import { ISelectionData, ExecutionPlanOptions } from 'azdata';
 import { UntitledEditorModel } from 'vs/workbench/common/editor/untitledEditorModel';
 import { IResolvedTextEditorModel } from 'vs/editor/common/services/resolverService';
-import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
-import QueryRunner from 'sql/platform/query/common/queryRunner';
-import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 
 const MAX_SIZE = 13;
 
@@ -112,56 +110,63 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 
 	private _updateSelection: Emitter<ISelectionData>;
 
-	private _runner: QueryRunner;
-	public get runner(): QueryRunner {
-		return this._runner;
-	}
-
-	private _results: QueryResultsInput;
-
 	constructor(
 		private _description: string,
 		private _sql: UntitledEditorInput,
+		private _results: QueryResultsInput,
 		private _connectionProviderName: string,
-		@IConnectionManagementService private readonly connectionManagementService: IConnectionManagementService,
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
-		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@IQueryModelService private _queryModelService: IQueryModelService,
-		@IInstantiationService instantiationService: IInstantiationService
+		@IConfigurationService private _configurationService: IConfigurationService
 	) {
 		super();
 		this._updateSelection = new Emitter<ISelectionData>();
 
-		this._runner = this._register(instantiationService.createInstance(QueryRunner, this.uri));
-		this._results = this._register(instantiationService.createInstance(QueryResultsInput, this.uri, this.runner));
-
 		this._register(this._sql);
+		this._register(this._results);
 
 		// re-emit sql editor events through this editor if it exists
 		if (this._sql) {
 			this._register(this._sql.onDidChangeDirty(() => this._onDidChangeDirty.fire()));
 		}
 
-		// Register callbacks for the Actions
+		// Attach to event callbacks
-		this.runner.onQueryStart(() => this.onRunQuery());
+		if (this._queryModelService) {
+			// Register callbacks for the Actions
+			this._register(
+				this._queryModelService.onRunQueryStart(uri => {
+					if (this.uri === uri) {
+						this.onRunQuery();
+					}
+				})
+			);
 
-		this.runner.onQueryEnd(() => this.onQueryComplete());
+			this._register(
+				this._queryModelService.onRunQueryComplete(uri => {
+					if (this.uri === uri) {
+						this.onQueryComplete();
+					}
+				})
+			);
+		}
 
-		if (this.connectionManagementService) {
+		if (this._connectionManagementService) {
-			this._register(this.connectionManagementService.onDisconnect(result => {
+			this._register(this._connectionManagementService.onDisconnect(result => {
 				if (result.connectionUri === this.uri) {
 					this.onDisconnect();
 				}
 			}));
 			if (this.uri) {
 				if (this._connectionProviderName) {
-					this.connectionManagementService.doChangeLanguageFlavor(this.uri, 'sql', this._connectionProviderName);
+					this._connectionManagementService.doChangeLanguageFlavor(this.uri, 'sql', this._connectionProviderName);
 				} else {
-					this.connectionManagementService.ensureDefaultLanguageFlavor(this.uri);
+					this._connectionManagementService.ensureDefaultLanguageFlavor(this.uri);
 				}
 			}
 		}
 
-		if (this.configurationService) {
+		if (this._configurationService) {
-			this._register(this.configurationService.onDidChangeConfiguration(e => {
+			this._register(this._configurationService.onDidChangeConfiguration(e => {
 				if (e.affectedKeys.includes('sql.showConnectionInfoInTitle')) {
 					this._onDidChangeLabel.fire();
 				}
@@ -214,8 +219,8 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 	}
 
 	public getName(longForm?: boolean): string {
-		if (this.configurationService.getValue('sql.showConnectionInfoInTitle')) {
+		if (this._configurationService.getValue('sql.showConnectionInfoInTitle')) {
-			let profile = this.connectionManagementService.getConnectionProfile(this.uri);
+			let profile = this._connectionManagementService.getConnectionProfile(this.uri);
 			let title = '';
 			if (this._description && this._description !== '') {
 				title = this._description + ' ';
@@ -248,17 +253,17 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 
 	// State update funtions
 	public runQuery(selection: ISelectionData, executePlanOptions?: ExecutionPlanOptions): void {
-		this.runner.runQuery(selection, executePlanOptions);
+		this._queryModelService.runQuery(this.uri, selection, this, executePlanOptions);
 		this.state.executing = true;
 	}
 
 	public runQueryStatement(selection: ISelectionData): void {
-		this.runner.runQueryStatement(selection);
+		this._queryModelService.runQueryStatement(this.uri, selection, this);
 		this.state.executing = true;
 	}
 
 	public runQueryString(text: string): void {
-		this.runner.runQuery(text);
+		this._queryModelService.runQueryString(this.uri, text, this);
 		this.state.executing = true;
 	}
 
@@ -281,7 +286,7 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 		this.state.connected = true;
 		this.state.connecting = false;
 
-		let isRunningQuery = this.runner.isExecuting;
+		let isRunningQuery = this._queryModelService.isRunningQuery(this.uri);
 		if (!isRunningQuery && params && params.runQueryOnCompletion) {
 			let selection: ISelectionData = params ? params.querySelection : undefined;
 			if (params.runQueryOnCompletion === RunQueryOnConnectionMode.executeCurrentQuery) {
@@ -312,7 +317,8 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 	}
 
 	public close(): void {
-		this.connectionManagementService.disconnectEditor(this, true);
+		this._queryModelService.disposeQuery(this.uri);
+		this._connectionManagementService.disconnectEditor(this, true);
 
 		this._sql.close();
 		this._results.close();
@@ -323,6 +329,6 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 	 * Get the color that should be displayed
 	 */
 	public get tabColor(): string {
-		return this.connectionManagementService.getTabColorForUri(this.uri);
+		return this._connectionManagementService.getTabColorForUri(this.uri);
 	}
 }

src/sql/workbench/parts/query/common/queryResultsInput.ts

@@ -13,7 +13,6 @@ import { QueryPlanState } from 'sql/workbench/parts/queryPlan/common/queryPlanSt
 import { MessagePanelState } from 'sql/workbench/parts/query/common/messagePanelState';
 import { GridPanelState } from 'sql/workbench/parts/query/common/gridPanelState';
 import { QueryModelViewState } from 'sql/workbench/parts/query/common/modelViewTab/modelViewState';
-import QueryRunner from 'sql/platform/query/common/queryRunner';
 
 export class ResultsViewState {
 	public gridPanelState: GridPanelState = new GridPanelState();
@@ -63,11 +62,7 @@ export class QueryResultsInput extends EditorInput {
 		return this._state;
 	}
 
-	public get runner(): QueryRunner {
+	constructor(private _uri: string) {
-		return this._runner;
-	}
-
-	constructor(private _uri: string, private _runner: QueryRunner) {
 		super();
 		this._visible = false;
 		this._hasBootstrapped = false;

src/sql/workbench/parts/query/test/browser/queryActions.test.ts

@@ -21,6 +21,7 @@ import {
 } from 'sql/workbench/parts/query/browser/queryActions';
 import { QueryInput } from 'sql/workbench/parts/query/common/queryInput';
 import { QueryEditor } from 'sql/workbench/parts/query/browser/queryEditor';
+import { QueryModelService } from 'sql/platform/query/common/queryModelService';
 import { ConnectionManagementService } from 'sql/platform/connection/common/connectionManagementService';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
@@ -30,11 +31,6 @@ import { TestStorageService, TestFileService } from 'vs/workbench/test/workbench
 import { MockContextKeyService } from 'vs/platform/keybinding/test/common/mockKeybindingService';
 import { TestThemeService } from 'vs/platform/theme/test/common/testThemeService';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
-import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
-import { TestEditorInput } from 'vs/workbench/services/editor/test/browser/editorGroupsService.test';
-import { URI } from 'vs/base/common/uri';
-import { Schemas } from 'vs/base/common/network';
 
 let none: void;
 
@@ -48,9 +44,7 @@ suite('SQL QueryAction Tests', () => {
 
 	setup(() => {
 		// Setup a reusable mock QueryInput
-		const instantiationService = new TestInstantiationService();
+		testQueryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
-		const fileInput = new TestEditorInput(URI.from({ scheme: Schemas.untitled }));
-		testQueryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict, undefined, fileInput, undefined, undefined, undefined, instantiationService);
 		testQueryInput.setup(x => x.uri).returns(() => testUri);
 		testQueryInput.setup(x => x.runQuery(undefined)).callback(() => { calledRunQueryOnInput = true; });
 
@@ -75,7 +69,7 @@ suite('SQL QueryAction Tests', () => {
 
 	test('setClass sets child CSS class correctly', (done) => {
 		// If I create a RunQueryAction
-		let queryAction: QueryTaskbarAction = new RunQueryAction(undefined, undefined);
+		let queryAction: QueryTaskbarAction = new RunQueryAction(undefined, undefined, undefined, undefined);
 
 		// "class should automatically get set to include the base class and the RunQueryAction class
 		let className = RunQueryAction.EnabledClass;
@@ -99,7 +93,7 @@ suite('SQL QueryAction Tests', () => {
 		editor.setup(x => x.input).returns(() => testQueryInput.object);
 
 		// If I create a QueryTaskbarAction and I pass a non-connected editor to _getConnectedQueryEditorUri
-		let queryAction: QueryTaskbarAction = new RunQueryAction(undefined, connectionManagementService.object);
+		let queryAction: QueryTaskbarAction = new RunQueryAction(undefined, undefined, connectionManagementService.object, undefined);
 		let connected: boolean = queryAction.isConnected(editor.object);
 
 		// I should get an unconnected state
@@ -135,8 +129,14 @@ suite('SQL QueryAction Tests', () => {
 		connectionManagementService.callBase = true;
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
 
+		// ... Mock QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+		queryModelService.setup(x => x.runQuery(TypeMoq.It.isAny(), undefined, TypeMoq.It.isAny(), TypeMoq.It.isAny())).callback(() => {
+			calledRunQuery = true;
+		});
+
 		// If I call run on RunQueryAction when I am not connected
-		let queryAction: RunQueryAction = new RunQueryAction(editor.object, connectionManagementService.object);
+		let queryAction: RunQueryAction = new RunQueryAction(editor.object, queryModelService.object, connectionManagementService.object, undefined);
 		isConnected = false;
 		calledRunQueryOnInput = false;
 		queryAction.run();
@@ -170,9 +170,8 @@ suite('SQL QueryAction Tests', () => {
 		let countCalledRunQuery: number = 0;
 
 		// ... Mock "isSelectionEmpty" in QueryEditor
-		const instantiationService = new TestInstantiationService();
+		let queryInput: TypeMoq.Mock<QueryInput> = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
-		const fileInput = new TestEditorInput(URI.from({ scheme: Schemas.untitled }));
+		queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
-		let queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict, undefined, fileInput, undefined, undefined, undefined, instantiationService);
 		queryInput.setup(x => x.uri).returns(() => testUri);
 		queryInput.setup(x => x.runQuery(undefined)).callback(() => {
 			countCalledRunQuery++;
@@ -192,8 +191,11 @@ suite('SQL QueryAction Tests', () => {
 		connectionManagementService.callBase = true;
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => true);
 
+		// ... Mock QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+
 		// If I call run on RunQueryAction when I have a non empty selection
-		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, connectionManagementService.object);
+		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, queryModelService.object, connectionManagementService.object, undefined);
 		isSelectionEmpty = false;
 		queryAction.run();
 
@@ -232,9 +234,7 @@ suite('SQL QueryAction Tests', () => {
 			.returns(() => Promise.resolve(none));
 
 		// ... Mock "getSelection" in QueryEditor
-		const instantiationService = new TestInstantiationService();
+		let queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Loose);
-		const fileInput = new TestEditorInput(URI.from({ scheme: Schemas.untitled }));
-		let queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Loose, undefined, fileInput, undefined, undefined, undefined, instantiationService);
 		queryInput.setup(x => x.uri).returns(() => testUri);
 		queryInput.setup(x => x.runQuery(TypeMoq.It.isAny())).callback((selection: ISelectionData) => {
 			runQuerySelection = selection;
@@ -266,7 +266,7 @@ suite('SQL QueryAction Tests', () => {
 		/// End Setup Test ///
 
 		////// If I call run on RunQueryAction while disconnected and with an undefined selection
-		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, connectionManagementService.object);
+		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, undefined, connectionManagementService.object, undefined);
 		isConnected = false;
 		selectionToReturnInGetSelection = undefined;
 		queryAction.run();
@@ -317,6 +317,38 @@ suite('SQL QueryAction Tests', () => {
 		assert.equal(runQuerySelection.endColumn, selectionToReturnInGetSelection.endColumn, 'endColumn should match');
 	});
 
+	test('CancelQueryAction calls cancelQuery() only if URI is connected', (done) => {
+		// ... Create assert variables
+		let isConnected: boolean = undefined;
+		let calledCancelQuery: boolean = false;
+
+		// ... Mock "isConnected" in ConnectionManagementService
+		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {});
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+
+		// ... Mock QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+		queryModelService.setup(x => x.cancelQuery(TypeMoq.It.isAny())).callback(() => {
+			calledCancelQuery = true;
+		});
+
+		// If I call run on CancelQueryAction when I am not connected
+		let queryAction: CancelQueryAction = new CancelQueryAction(editor.object, queryModelService.object, connectionManagementService.object);
+		isConnected = false;
+		queryAction.run();
+
+		// cancelQuery should not be run
+		assert.equal(calledCancelQuery, false, 'run should not call cancelQuery');
+
+		// If I call run on CancelQueryAction when I am connected
+		isConnected = true;
+		queryAction.run();
+
+		// cancelQuery should be run
+		assert.equal(calledCancelQuery, true, 'run should call cancelQuery');
+		done();
+	});
+
 	// We want to call disconnectEditor regardless of connection to be able to cancel in-progress connections
 	test('DisconnectDatabaseAction calls disconnectEditor regardless of URI being connected', (done) => {
 		// ... Create assert variables

src/sql/workbench/parts/query/test/browser/queryEditor.test.ts

@@ -10,6 +10,7 @@ import { Memento } from 'vs/workbench/common/memento';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
 
 import { QueryResultsInput } from 'sql/workbench/parts/query/common/queryResultsInput';
+import { QueryModelService } from 'sql/platform/query/common/queryModelService';
 import { QueryInput } from 'sql/workbench/parts/query/common/queryInput';
 import { INewConnectionParams, ConnectionType, RunQueryOnConnectionMode } from 'sql/platform/connection/common/connectionManagement';
 import { ConnectionManagementService } from 'sql/platform/connection/common/connectionManagementService';
@@ -22,8 +23,6 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
 import { BaseEditor } from 'vs/workbench/browser/parts/editor/baseEditor';
 import { TestStorageService } from 'vs/workbench/test/workbenchTestServices';
 import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';
-import QueryRunner from 'sql/platform/query/common/queryRunner';
-import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
 
 suite('SQL QueryEditor Tests', () => {
 	let instantiationService: TypeMoq.Mock<InstantiationService>;
@@ -55,7 +54,7 @@ suite('SQL QueryEditor Tests', () => {
 			return new Promise((resolve) => resolve(mockEditor));
 		});
 		instantiationService.setup(x => x.createInstance(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((input) => {
-			return new Promise((resolve) => resolve(new RunQueryAction(undefined, undefined)));
+			return new Promise((resolve) => resolve(new RunQueryAction(undefined, undefined, undefined, undefined)));
 		});
 		// Setup hook to capture calls to create the listDatabase action
 		instantiationService.setup(x => x.createInstance(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((classDef, editor, action) => {
@@ -65,7 +64,7 @@ suite('SQL QueryEditor Tests', () => {
 				}
 			}
 			// Default
-			return new RunQueryAction(undefined, undefined);
+			return new RunQueryAction(undefined, undefined, undefined, undefined);
 		});
 
 		// Mock EditorDescriptorService to give us a mock editor description
@@ -248,6 +247,7 @@ suite('SQL QueryEditor Tests', () => {
 	suite('Action Tests', () => {
 		let queryActionInstantiationService: TypeMoq.Mock<InstantiationService>;
 		let queryConnectionService: TypeMoq.Mock<ConnectionManagementService>;
+		let queryModelService: TypeMoq.Mock<QueryModelService>;
 		let queryInput: QueryInput;
 		setup(() => {
 
@@ -262,7 +262,6 @@ suite('SQL QueryEditor Tests', () => {
 
 			// Mock InstantiationService to give us the actions
 			queryActionInstantiationService = TypeMoq.Mock.ofType(InstantiationService, TypeMoq.MockBehavior.Loose);
-			instantiationService.setup(x => x.createInstance(TypeMoq.It.isValue(QueryRunner), TypeMoq.It.isAnyString())).returns(() => new QueryRunner('', undefined, undefined, undefined, undefined, undefined));
 
 			queryActionInstantiationService.setup(x => x.createInstance(TypeMoq.It.isAny())).returns((input) => {
 				return new Promise((resolve) => resolve(mockEditor));
@@ -270,7 +269,7 @@ suite('SQL QueryEditor Tests', () => {
 
 			queryActionInstantiationService.setup(x => x.createInstance(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((input) => {
 				// Default
-				return new RunQueryAction(undefined, undefined);
+				return new RunQueryAction(undefined, undefined, undefined, undefined);
 			});
 
 			// Setup hook to capture calls to create the listDatabase action
@@ -281,24 +280,27 @@ suite('SQL QueryEditor Tests', () => {
 						return item;
 					}
 					// Default
-					return new RunQueryAction(undefined, undefined);
+					return new RunQueryAction(undefined, undefined, undefined, undefined);
 				});
 
-			const mockInstantiationService = new TestInstantiationService();
-
 			let fileInput = new UntitledEditorInput(URI.parse('file://testUri'), false, '', '', '', instantiationService.object, undefined, undefined);
+			queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose, undefined, undefined);
+			queryModelService.callBase = true;
+			queryModelService.setup(x => x.disposeQuery(TypeMoq.It.isAny())).returns(() => void 0);
 			queryInput = new QueryInput(
 				'',
 				fileInput,
 				undefined,
-				connectionManagementService.object,
 				undefined,
-				mockInstantiationService
+				connectionManagementService.object,
+				queryModelService.object,
+				undefined
 			);
 		});
 
 		test('Taskbar buttons are set correctly upon standard load', () => {
 			queryConnectionService.setup(x => x.isConnected(TypeMoq.It.isAny())).returns(() => false);
+			queryModelService.setup(x => x.isRunningQuery(TypeMoq.It.isAny())).returns(() => false);
 			// If I use the created QueryEditor with no changes since creation
 			// Buttons should be set as if disconnected
 			assert.equal(queryInput.state.connected, false, 'query state should be not connected');
@@ -309,14 +311,16 @@ suite('SQL QueryEditor Tests', () => {
 		test('Taskbar buttons are set correctly upon connect', () => {
 			let params: INewConnectionParams = { connectionType: ConnectionType.editor, runQueryOnCompletion: RunQueryOnConnectionMode.none };
 			queryInput.onConnectSuccess(params);
+			queryModelService.setup(x => x.isRunningQuery(TypeMoq.It.isAny())).returns(() => false);
 			assert.equal(queryInput.state.connected, true, 'query state should be not connected');
 			assert.equal(queryInput.state.executing, false, 'query state should be not executing');
 			assert.equal(queryInput.state.connecting, false, 'query state should be not connecting');
 		});
 		test('Test that we attempt to dispose query when the queryInput is disposed', () => {
-			let queryResultsInput = new QueryResultsInput('testUri', undefined);
+			let queryResultsInput = new QueryResultsInput('testUri');
 			queryInput['_results'] = queryResultsInput;
 			queryInput.close();
+			queryModelService.verify(x => x.disposeQuery(TypeMoq.It.isAnyString()), TypeMoq.Times.once());
 		});
 	});
 });

src/sql/workbench/services/queryEditor/browser/queryEditorService.ts

@@ -12,6 +12,7 @@ import { QueryPlanInput } from 'sql/workbench/parts/queryPlan/common/queryPlanIn
 import { sqlModeId, untitledFilePrefix, getSupportedInputResource } from 'sql/workbench/common/customInputConverter';
 import * as TaskUtilities from 'sql/workbench/browser/taskUtilities';
 
+import { IMode } from 'vs/editor/common/modes';
 import { ITextModel } from 'vs/editor/common/model';
 import { IUntitledEditorService } from 'vs/workbench/services/untitled/common/untitledEditorService';
 import { IEditorService, ACTIVE_GROUP } from 'vs/workbench/services/editor/common/editorService';
@@ -90,9 +91,10 @@ export class QueryEditorService implements IQueryEditorService {
 					}
 				}
 
-				let input = this._instantiationService.createInstance(QueryInput, objectName, fileInput, connectionProviderName);
+				const queryResultsInput: QueryResultsInput = this._instantiationService.createInstance(QueryResultsInput, docUri.toString());
+				let queryInput: QueryInput = this._instantiationService.createInstance(QueryInput, objectName, fileInput, queryResultsInput, connectionProviderName);
 
-				this._editorService.openEditor(input, { pinned: true })
+				this._editorService.openEditor(queryInput, { pinned: true })
 					.then((editor) => {
 						let params = <QueryInput>editor.input;
 						resolve(params);
@@ -286,7 +288,8 @@ export class QueryEditorService implements IQueryEditorService {
 
 		let newEditorInput: IEditorInput = undefined;
 		if (changingToSql) {
-			let queryInput: QueryInput = QueryEditorService.instantiationService.createInstance(QueryInput, '', input, undefined);
+			const queryResultsInput: QueryResultsInput = QueryEditorService.instantiationService.createInstance(QueryResultsInput, uri.toString());
+			let queryInput: QueryInput = QueryEditorService.instantiationService.createInstance(QueryInput, '', input, queryResultsInput, undefined);
 			newEditorInput = queryInput;
 		} else {
 			let uriCopy: URI = URI.from({ scheme: uri.scheme, authority: uri.authority, path: uri.path, query: uri.query, fragment: uri.fragment });

src/sql/workbench/test/common/taskUtilities.test.ts

@@ -14,7 +14,6 @@ import { URI } from 'vs/base/common/uri';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
 import { QueryInput } from 'sql/workbench/parts/query/common/queryInput';
 import { TestEditorService } from 'vs/workbench/test/workbenchTestServices';
-import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
 
 suite('TaskUtilities', function () {
 	test('getCurrentGlobalConnection returns the selected OE server if a server or one of its children is selected', () => {
@@ -77,8 +76,7 @@ suite('TaskUtilities', function () {
 		// Mock the workbench service to return the active tab connection
 		let tabConnectionUri = 'file://test_uri';
 		let editorInput = new UntitledEditorInput(URI.parse(tabConnectionUri), false, undefined, undefined, undefined, undefined, undefined, undefined);
-		const instantiationService = new TestInstantiationService();
+		let queryInput = new QueryInput(undefined, editorInput, undefined, undefined, undefined, undefined, undefined);
-		let queryInput = new QueryInput(undefined, editorInput, undefined, undefined, undefined, instantiationService);
 		mockConnectionManagementService.setup(x => x.getConnectionProfile(tabConnectionUri)).returns(() => tabProfile);
 
 		// If I call getCurrentGlobalConnection, it should return the expected profile from the active tab