Rework how we handle custom editors (#5696)

* update how we handle editors * small edit * handle changing languages * implement generic language association * implement notebook serializers * fix tests * formatting * update how we handle editors * small edit * handle changing languages * implement generic language association * implement notebook serializers * fix tests * formatting * fix broken * fix compile * fix tests * add back in removed note book contributions * fix layering * fix compile errors * fix workbench * fix hanging promises * idk why these changed * fix change * add comments to language change code * fix a few bugs * add query plan association
2026-01-25 09:35:37 -05:00 · 2019-11-24 19:22:11 -08:00
parent f3a6fc6f88
commit 43387f0d0b
50 changed files with 988 additions and 873 deletions
--- a/src/sql/workbench/contrib/query/browser/keyboardQueryActions.ts
+++ b/src/sql/workbench/contrib/query/browser/keyboardQueryActions.ts
@@ -20,7 +20,7 @@ import * as ConnectionConstants from 'sql/platform/connection/common/constants';
 import { EditDataEditor } from 'sql/workbench/contrib/editData/browser/editDataEditor';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-import { QueryInput } from 'sql/workbench/contrib/query/common/queryInput';
+import { QueryEditorInput } from 'sql/workbench/contrib/query/common/queryEditorInput';
 import { firstIndex } from 'vs/base/common/arrays';

 const singleQuote = '\'';
@@ -199,7 +199,7 @@ export class RefreshIntellisenseKeyboardAction extends Action {

 	public run(): Promise<void> {
 		const editor = this.editorService.activeEditor;
-		if (editor instanceof QueryInput) {
+		if (editor instanceof QueryEditorInput) {
 			this.connectionManagementService.rebuildIntelliSenseCache(editor.uri);
 		}
 		return Promise.resolve(null);
--- a/src/sql/workbench/contrib/query/browser/query.contribution.ts
+++ b/src/sql/workbench/contrib/query/browser/query.contribution.ts
@@ -19,7 +19,6 @@ import { QueryEditor } from 'sql/workbench/contrib/query/browser/queryEditor';
 import { QueryResultsEditor } from 'sql/workbench/contrib/query/browser/queryResultsEditor';
 import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResultsInput';
 import * as queryContext from 'sql/workbench/contrib/query/common/queryContext';
-import { QueryInput } from 'sql/workbench/contrib/query/common/queryInput';
 import {
 	RunQueryKeyboardAction, RunCurrentQueryKeyboardAction, CancelQueryKeyboardAction, RefreshIntellisenseKeyboardAction, ToggleQueryResultsKeyboardAction,
 	RunQueryShortcutAction, RunCurrentQueryWithActualPlanKeyboardAction, FocusOnCurrentQueryKeyboardAction, ParseSyntaxAction
@@ -33,6 +32,14 @@ import { IWorkbenchContributionsRegistry, Extensions as WorkbenchExtensions } fr
 import { LifecyclePhase } from 'vs/platform/lifecycle/common/lifecycle';
 import { TimeElapsedStatusBarContributions, RowCountStatusBarContributions, QueryStatusStatusBarContributions } from 'sql/workbench/contrib/query/browser/statusBarItems';
 import { SqlFlavorStatusbarItem } from 'sql/workbench/contrib/query/browser/flavorStatus';
+import { IEditorInputFactoryRegistry, Extensions as EditorInputFactoryExtensions } from 'vs/workbench/common/editor';
+import { FileQueryEditorInput } from 'sql/workbench/contrib/query/common/fileQueryEditorInput';
+import { FileQueryEditorInputFactory, UntitledQueryEditorInputFactory } from 'sql/workbench/contrib/query/common/queryInputFactory';
+import { UntitledQueryEditorInput } from 'sql/workbench/contrib/query/common/untitledQueryEditorInput';
+import { ILanguageAssociationRegistry, Extensions as LanguageAssociationExtensions } from 'sql/workbench/common/languageAssociation';
+import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileEditorInput';
+import { QueryEditorInput } from 'sql/workbench/contrib/query/common/queryEditorInput';
+import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
 import { NewQueryTask, OE_NEW_QUERY_ACTION_ID, DE_NEW_QUERY_COMMAND_ID } from 'sql/workbench/contrib/query/browser/queryActions';
 import { TreeNodeContextKey } from 'sql/workbench/contrib/objectExplorer/common/treeNodeContextKey';
 import { MssqlNodeContext } from 'sql/workbench/contrib/dataExplorer/browser/mssqlNodeContext';
@@ -44,25 +51,30 @@ export const QueryEditorVisibleCondition = ContextKeyExpr.has(queryContext.query
 export const ResultsGridFocusCondition = ContextKeyExpr.and(ContextKeyExpr.has(queryContext.resultsVisibleId), ContextKeyExpr.has(queryContext.resultsGridFocussedId));
 export const ResultsMessagesFocusCondition = ContextKeyExpr.and(ContextKeyExpr.has(queryContext.resultsVisibleId), ContextKeyExpr.has(queryContext.resultsMessagesFocussedId));

-// Editor
-const queryResultsEditorDescriptor = new EditorDescriptor(
-	QueryResultsEditor,
-	QueryResultsEditor.ID,
-	'QueryResults'
-);
+Registry.as<IEditorInputFactoryRegistry>(EditorInputFactoryExtensions.EditorInputFactories)
+	.registerEditorInputFactory(FileQueryEditorInput.ID, FileQueryEditorInputFactory);
+
+Registry.as<IEditorInputFactoryRegistry>(EditorInputFactoryExtensions.EditorInputFactories)
+	.registerEditorInputFactory(UntitledQueryEditorInput.ID, UntitledQueryEditorInputFactory);
+
+Registry.as<ILanguageAssociationRegistry>(LanguageAssociationExtensions.LanguageAssociations)
+	.registerLanguageAssociation('sql', (accessor, editor) => {
+		const instantiationService = accessor.get(IInstantiationService);
+		const queryResultsInput = instantiationService.createInstance(QueryResultsInput, editor.getResource().toString());
+		if (editor instanceof FileEditorInput) {
+			return instantiationService.createInstance(FileQueryEditorInput, '', editor, queryResultsInput);
+		} else if (editor instanceof UntitledEditorInput) {
+			return instantiationService.createInstance(UntitledQueryEditorInput, '', editor, queryResultsInput);
+		} else {
+			return undefined;
+		}
+	}, (editor: QueryEditorInput) => editor.text, true);

 Registry.as<IEditorRegistry>(EditorExtensions.Editors)
-	.registerEditor(queryResultsEditorDescriptor, [new SyncDescriptor(QueryResultsInput)]);
-
-// Editor
-const queryEditorDescriptor = new EditorDescriptor(
-	QueryEditor,
-	QueryEditor.ID,
-	'Query'
-);
+	.registerEditor(new EditorDescriptor(QueryResultsEditor, QueryResultsEditor.ID, localize('queryResultsEditor.name', "Query Results")), [new SyncDescriptor(QueryResultsInput)]);

 Registry.as<IEditorRegistry>(EditorExtensions.Editors)
-	.registerEditor(queryEditorDescriptor, [new SyncDescriptor(QueryInput)]);
+	.registerEditor(new EditorDescriptor(QueryEditor, QueryEditor.ID, localize('queryEditor.name', "Query Editor")), [new SyncDescriptor(FileQueryEditorInput), new SyncDescriptor(UntitledQueryEditorInput)]);

 const actionRegistry = <IWorkbenchActionRegistry>Registry.as(Extensions.WorkbenchActions);

--- a/src/sql/workbench/contrib/query/browser/queryEditor.ts
+++ b/src/sql/workbench/contrib/query/browser/queryEditor.ts
@@ -31,7 +31,7 @@ import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileE
 import { URI } from 'vs/base/common/uri';
 import { IFileService, FileChangesEvent } from 'vs/platform/files/common/files';

-import { QueryInput, IQueryEditorStateChange } from 'sql/workbench/contrib/query/common/queryInput';
+import { QueryEditorInput, IQueryEditorStateChange } from 'sql/workbench/contrib/query/common/queryEditorInput';
 import { QueryResultsEditor } from 'sql/workbench/contrib/query/browser/queryResultsEditor';
 import * as queryContext from 'sql/workbench/contrib/query/common/queryContext';
 import { Taskbar, ITaskbarContent } from 'sql/base/browser/ui/taskbar/taskbar';
@@ -120,8 +120,8 @@ export class QueryEditor extends BaseEditor {
 	}

 	// PUBLIC METHODS ////////////////////////////////////////////////////////////
-	public get input(): QueryInput | null {
-		return this._input as QueryInput;
+	public get input(): QueryEditorInput | null {
+		return this._input as QueryEditorInput;
 	}

 	/**
@@ -278,7 +278,7 @@ export class QueryEditor extends BaseEditor {
 		this.taskbar.setContent(content);
 	}

-	public async setInput(newInput: QueryInput, options: EditorOptions, token: CancellationToken): Promise<void> {
+	public async setInput(newInput: QueryEditorInput, options: EditorOptions, token: CancellationToken): Promise<void> {
 		const oldInput = this.input;

 		if (newInput.matches(oldInput)) {
@@ -293,7 +293,7 @@ export class QueryEditor extends BaseEditor {
 		}

 		// If we're switching editor types switch out the views
-		const newTextEditor = newInput.sql instanceof FileEditorInput ? this.textFileEditor : this.textResourceEditor;
+		const newTextEditor = newInput.text instanceof FileEditorInput ? this.textFileEditor : this.textResourceEditor;
 		if (newTextEditor !== this.currentTextEditor) {
 			this.currentTextEditor = newTextEditor;
 			this.splitview.removeView(0, Sizing.Distribute);
@@ -309,7 +309,7 @@ export class QueryEditor extends BaseEditor {

 		await Promise.all([
 			super.setInput(newInput, options, token),
-			this.currentTextEditor.setInput(newInput.sql, options, token),
+			this.currentTextEditor.setInput(newInput.text, options, token),
 			this.resultsEditor.setInput(newInput.results, options)
 		]);

@@ -324,7 +324,7 @@ export class QueryEditor extends BaseEditor {
 		}
 	}

-	private saveQueryEditorViewState(input: QueryInput): void {
+	private saveQueryEditorViewState(input: QueryEditorInput): void {
 		if (!input) {
 			return; // ensure we have an input to handle view state for
 		}
--- a/src/sql/workbench/contrib/query/browser/queryResultsEditor.ts
+++ b/src/sql/workbench/contrib/query/browser/queryResultsEditor.ts
@@ -19,7 +19,7 @@ import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResul
 import { QueryResultsView } from 'sql/workbench/contrib/query/browser/queryResultsView';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { IStorageService } from 'vs/platform/storage/common/storage';
-import { RESULTS_GRID_DEFAULTS } from 'sql/workbench/contrib/query/common/resultsGridContribution';
+import { RESULTS_GRID_DEFAULTS } from 'sql/workbench/contrib/query/common/resultsGrid.contribution';

 export const TextCompareEditorVisible = new RawContextKey<boolean>('textCompareEditorVisible', false);

--- a/src/sql/workbench/contrib/query/browser/statusBarItems.ts
+++ b/src/sql/workbench/contrib/query/browser/statusBarItems.ts
@@ -9,10 +9,10 @@ import { IntervalTimer } from 'vs/base/common/async';
 import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { IWorkbenchContribution } from 'vs/workbench/common/contributions';
 import { localize } from 'vs/nls';
-import { QueryInput } from 'sql/workbench/contrib/query/common/queryInput';
 import QueryRunner from 'sql/platform/query/common/queryRunner';
 import { parseNumAsTimeString } from 'sql/platform/connection/common/utils';
 import { Event } from 'vs/base/common/event';
+import { QueryEditorInput } from 'sql/workbench/contrib/query/common/queryEditorInput';
 import { IStatusbarService, IStatusbarEntryAccessor, StatusbarAlignment } from 'vs/workbench/services/statusbar/common/statusbar';

 export class TimeElapsedStatusBarContributions extends Disposable implements IWorkbenchContribution {
@@ -56,7 +56,7 @@ export class TimeElapsedStatusBarContributions extends Disposable implements IWo
 		this.disposable.clear();
 		this.hide();
 		const activeInput = this.editorService.activeEditor;
-		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
+		if (activeInput && activeInput instanceof QueryEditorInput && activeInput.uri) {
 			const uri = activeInput.uri;
 			const runner = this.queryModelService.getQueryRunner(uri);
 			if (runner) {
@@ -148,7 +148,7 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 		this.disposable.clear();
 		this.hide();
 		const activeInput = this.editorService.activeEditor;
-		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
+		if (activeInput && activeInput instanceof QueryEditorInput && activeInput.uri) {
 			const uri = activeInput.uri;
 			const runner = this.queryModelService.getQueryRunner(uri);
 			if (runner) {
@@ -227,7 +227,7 @@ export class QueryStatusStatusBarContributions extends Disposable implements IWo
 		this.hide();
 		this.visisbleUri = undefined;
 		const activeInput = this.editorService.activeEditor;
-		if (activeInput && activeInput instanceof QueryInput && activeInput.uri) {
+		if (activeInput && activeInput instanceof QueryEditorInput && activeInput.uri) {
 			this.visisbleUri = activeInput.uri;
 			const runner = this.queryModelService.getQueryRunner(this.visisbleUri);
 			if (runner && runner.isExecuting) {
--- a/src/sql/workbench/contrib/query/common/fileQueryEditorInput.ts
+++ b/src/sql/workbench/contrib/query/common/fileQueryEditorInput.ts
@@ -0,0 +1,87 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { QueryEditorInput } from 'sql/workbench/contrib/query/common/queryEditorInput';
+import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResultsInput';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
+
+import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileEditorInput';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { EncodingMode } from 'vs/workbench/common/editor';
+import { TextFileEditorModel } from 'vs/workbench/services/textfile/common/textFileEditorModel';
+import { BinaryEditorModel } from 'vs/workbench/common/editor/binaryEditorModel';
+import { IFileService } from 'vs/platform/files/common/files';
+
+type PublicPart<T> = { [K in keyof T]: T[K] };
+
+export class FileQueryEditorInput extends QueryEditorInput implements PublicPart<FileEditorInput> {
+
+	public static readonly ID = 'workbench.editorInput.fileQueryInput';
+
+	constructor(
+		description: string,
+		text: FileEditorInput,
+		results: QueryResultsInput,
+		@IConnectionManagementService connectionManagementService: IConnectionManagementService,
+		@IQueryModelService queryModelService: IQueryModelService,
+		@IConfigurationService configurationService: IConfigurationService,
+		@IFileService fileService: IFileService
+	) {
+		super(description, text, results, connectionManagementService, queryModelService, configurationService, fileService);
+	}
+
+	public resolve(): Promise<TextFileEditorModel | BinaryEditorModel> {
+		return this.text.resolve();
+	}
+
+	public get text(): FileEditorInput {
+		return this._text as FileEditorInput;
+	}
+
+	public getTypeId(): string {
+		return FileQueryEditorInput.ID;
+	}
+
+	public getEncoding(): string {
+		return this.text.getEncoding();
+	}
+
+	public setEncoding(encoding: string, mode: EncodingMode) {
+		this.text.setEncoding(encoding, mode);
+	}
+
+	public getPreferredEncoding(): string {
+		return this.text.getPreferredEncoding();
+	}
+
+	public setPreferredEncoding(encoding: string) {
+		this.text.setPreferredEncoding(encoding);
+	}
+
+	public getPreferredMode(): string {
+		return this.text.getPreferredMode();
+	}
+
+	public setMode(mode: string) {
+		this.text.setMode(mode);
+	}
+
+	public setPreferredMode(mode: string) {
+		this.text.setPreferredMode(mode);
+	}
+
+	public setForceOpenAsText() {
+		this.text.setForceOpenAsText();
+	}
+
+	public setForceOpenAsBinary() {
+		this.text.setForceOpenAsBinary();
+	}
+
+	public isResolved(): boolean {
+		return this.text.isResolved();
+	}
+}
--- a/src/sql/workbench/contrib/query/common/queryEditorInput.ts
+++ b/src/sql/workbench/contrib/query/common/queryEditorInput.ts
@@ -5,10 +5,9 @@

 import { localize } from 'vs/nls';
 import { IDisposable, Disposable } from 'vs/base/common/lifecycle';
-import { Event, Emitter } from 'vs/base/common/event';
+import { Emitter } from 'vs/base/common/event';
 import { URI } from 'vs/base/common/uri';
-import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
-import { EditorInput, ConfirmResult, EncodingMode, IEncodingSupport } from 'vs/workbench/common/editor';
+import { EditorInput, ConfirmResult } from 'vs/workbench/common/editor';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IFileService } from 'vs/platform/files/common/files';

@@ -17,14 +16,10 @@ import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResul
 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 { startsWith } from 'vs/base/common/strings';

 const MAX_SIZE = 13;

-type PublicPart<T> = { [K in keyof T]: T[K] };
-
 function trimTitle(title: string): string {
 	const length = title.length;
 	const diff = length - MAX_SIZE;
@@ -115,79 +110,48 @@ export class QueryEditorState extends Disposable {
 * Input for the QueryEditor. This input is simply a wrapper around a QueryResultsInput for the QueryResultsEditor
 * and a UntitledEditorInput for the SQL File Editor.
 */
-export class QueryInput extends EditorInput implements IEncodingSupport, IConnectableInput, PublicPart<UntitledEditorInput>, IDisposable {
+export abstract class QueryEditorInput extends EditorInput implements IConnectableInput, IDisposable {

-	public static ID: string = 'workbench.editorinputs.queryInput';
 	public static SCHEMA: string = 'sql';

 	private _state = this._register(new QueryEditorState());
 	public get state(): QueryEditorState { return this._state; }

-	private _updateSelection: Emitter<ISelectionData>;
-
 	constructor(
 		private _description: string,
-		private _sql: UntitledEditorInput,
-		private _results: QueryResultsInput,
-		private _connectionProviderName: string,
-		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
-		@IQueryModelService private _queryModelService: IQueryModelService,
-		@IConfigurationService private _configurationService: IConfigurationService,
+		protected _text: EditorInput,
+		protected _results: QueryResultsInput,
+		@IConnectionManagementService private readonly connectionManagementService: IConnectionManagementService,
+		@IQueryModelService private readonly queryModelService: IQueryModelService,
+		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IFileService private _fileService: IFileService
 	) {
 		super();
-		this._updateSelection = new Emitter<ISelectionData>();

-		this._register(this._sql);
+		this._register(this._text);
 		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()));
-		}
+		this._text.onDidChangeDirty(() => this._onDidChangeDirty.fire());

-		// Attach to event callbacks
-		if (this._queryModelService) {
-			// Register callbacks for the Actions
-			this._register(
-				this._queryModelService.onRunQueryStart(uri => {
-					if (this.uri === uri) {
-						this.onRunQuery();
-					}
-				})
-			);
-
-			this._register(
-				this._queryModelService.onRunQueryComplete(uri => {
-					if (this.uri === uri) {
-						this.onQueryComplete();
-					}
-				})
-			);
-		}
-
-		if (this._connectionManagementService) {
-			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);
-				} else {
-					this._connectionManagementService.ensureDefaultLanguageFlavor(this.uri);
+		this._register(
+			this.queryModelService.onRunQueryComplete(uri => {
+				if (this.uri === uri) {
+					this.onQueryComplete();
 				}
+			})
+		);
+
+		this._register(this.connectionManagementService.onDisconnect(result => {
+			if (result.connectionUri === this.uri) {
+				this.onDisconnect();
 			}
-		}
+		}));

-		if (this._configurationService) {
-			this._register(this._configurationService.onDidChangeConfiguration(e => {
-				if (e.affectsConfiguration('sql.showConnectionInfoInTitle')) {
-					this._onDidChangeLabel.fire();
-				}
-			}));
-		}
+		this._register(this.configurationService.onDidChangeConfiguration(e => {
+			if (e.affectedKeys.indexOf('sql.showConnectionInfoInTitle') > -1) {
+				this._onDidChangeLabel.fire();
+			}
+		}));

 		this.onDisconnect();
 		this.onQueryComplete();
@@ -195,47 +159,30 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec

 	// Getters for private properties
 	public get uri(): string { return this.getResource().toString(true); }
-	public get sql(): UntitledEditorInput { return this._sql; }
+	public get text(): EditorInput { return this._text; }
 	public get results(): QueryResultsInput { return this._results; }
-	public updateSelection(selection: ISelectionData): void { this._updateSelection.fire(selection); }
-	public getTypeId(): string { return QueryInput.ID; }
 	// Description is shown beside the tab name in the combobox of open editors
 	public getDescription(): string { return this._description; }
 	public supportsSplitEditor(): boolean { return false; }
-	public getMode(): string { return QueryInput.SCHEMA; }
-	public revert(): Promise<boolean> { return this._sql.revert(); }
-	public setMode(mode: string) {
-		this._sql.setMode(mode);
-	}
+	public revert(): Promise<boolean> { return this._text.revert(); }

 	public matches(otherInput: any): boolean {
-		if (otherInput instanceof QueryInput) {
-			return this._sql.matches(otherInput.sql);
+		// we want to be able to match against our underlying input as well, bascially we are our underlying input
+		if (otherInput instanceof QueryEditorInput) {
+			return this._text.matches(otherInput._text);
+		} else {
+			return this._text.matches(otherInput);
 		}
-
-		return this._sql.matches(otherInput);
 	}

 	// Forwarding resource functions to the inline sql file editor
-	public get onDidModelChangeContent(): Event<void> { return this._sql.onDidModelChangeContent; }
-	public get onDidModelChangeEncoding(): Event<void> { return this._sql.onDidModelChangeEncoding; }
-	public resolve(): Promise<UntitledEditorModel & IResolvedTextEditorModel> { return this._sql.resolve(); }
-	public save(): Promise<boolean> { return this._sql.save(); }
-	public isDirty(): boolean { return this._sql.isDirty(); }
-	public confirmSave(): Promise<ConfirmResult> { return this._sql.confirmSave(); }
-	public getResource(): URI { return this._sql.getResource(); }
-	public getEncoding(): string { return this._sql.getEncoding(); }
-	public suggestFileName(): string { return this._sql.suggestFileName(); }
-	hasBackup(): boolean {
-		if (this.sql) {
-			return this.sql.hasBackup();
-		}
-
-		return false;
-	}
+	public save(): Promise<boolean> { return this._text.save(); }
+	public isDirty(): boolean { return this._text.isDirty(); }
+	public confirmSave(): Promise<ConfirmResult> { return this._text.confirmSave(); }
+	public getResource(): URI { return this._text.getResource(); }

 	public matchInputInstanceType(inputType: any): boolean {
-		return (this._sql instanceof inputType);
+		return (this._text instanceof inputType);
 	}

 	public inputFileExists(): Promise<boolean> {
@@ -243,8 +190,8 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 	}

 	public getName(longForm?: boolean): string {
-		if (this._configurationService.getValue('sql.showConnectionInfoInTitle')) {
-			let profile = this._connectionManagementService.getConnectionProfile(this.uri);
+		if (this.configurationService.getValue('sql.showConnectionInfoInTitle')) {
+			let profile = this.connectionManagementService.getConnectionProfile(this.uri);
 			let title = '';
 			if (this._description && this._description !== '') {
 				title = this._description + ' ';
@@ -258,36 +205,30 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 			} else {
 				title += localize('disconnected', "disconnected");
 			}
-			return this._sql.getName() + (longForm ? (' - ' + title) : ` - ${trimTitle(title)}`);
+			return this._text.getName() + (longForm ? (' - ' + title) : ` - ${trimTitle(title)}`);
 		} else {
-			return this._sql.getName();
+			return this._text.getName();
 		}
 	}

 	// Called to get the tooltip of the tab
-	public getTitle() {
+	public getTitle(): string {
 		return this.getName(true);
 	}

-	public get hasAssociatedFilePath(): boolean { return this._sql.hasAssociatedFilePath; }
-
-	public setEncoding(encoding: string, mode: EncodingMode /* ignored, we only have Encode */): void {
-		this._sql.setEncoding(encoding, mode);
-	}
-
 	// State update funtions
 	public runQuery(selection?: ISelectionData, executePlanOptions?: ExecutionPlanOptions): void {
-		this._queryModelService.runQuery(this.uri, selection, this, executePlanOptions);
+		this.queryModelService.runQuery(this.uri, selection, this, executePlanOptions);
 		this.state.executing = true;
 	}

 	public runQueryStatement(selection?: ISelectionData): void {
-		this._queryModelService.runQueryStatement(this.uri, selection, this);
+		this.queryModelService.runQueryStatement(this.uri, selection, this);
 		this.state.executing = true;
 	}

 	public runQueryString(text: string): void {
-		this._queryModelService.runQueryString(this.uri, text, this);
+		this.queryModelService.runQueryString(this.uri, text, this);
 		this.state.executing = true;
 	}

@@ -314,7 +255,7 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 		this.state.connected = true;
 		this.state.connecting = false;

-		let isRunningQuery = this._queryModelService.isRunningQuery(this.uri);
+		let isRunningQuery = this.queryModelService.isRunningQuery(this.uri);
 		if (!isRunningQuery && params && params.runQueryOnCompletion) {
 			let selection: ISelectionData | undefined = params ? params.querySelection : undefined;
 			if (params.runQueryOnCompletion === RunQueryOnConnectionMode.executeCurrentQuery) {
@@ -345,10 +286,10 @@ export class QueryInput extends EditorInput implements IEncodingSupport, IConnec
 	}

 	public close(): void {
-		this._queryModelService.disposeQuery(this.uri);
-		this._connectionManagementService.disconnectEditor(this, true);
+		this.queryModelService.disposeQuery(this.uri);
+		this.connectionManagementService.disconnectEditor(this, true);

-		this._sql.close();
+		this._text.close();
 		this._results.close();
 		super.close();
 	}
@@ -357,7 +298,7 @@ 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);
 	}

 	public get isSharedSession(): boolean {
--- a/src/sql/workbench/contrib/query/common/queryInputFactory.ts
+++ b/src/sql/workbench/contrib/query/common/queryInputFactory.ts
@@ -0,0 +1,50 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IEditorInputFactory, IEditorInputFactoryRegistry, Extensions as EditorInputExtensions } from 'vs/workbench/common/editor';
+import { Registry } from 'vs/platform/registry/common/platform';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
+import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResultsInput';
+import { FILE_EDITOR_INPUT_ID } from 'vs/workbench/contrib/files/common/files';
+import { UntitledQueryEditorInput } from 'sql/workbench/contrib/query/common/untitledQueryEditorInput';
+import { FileQueryEditorInput } from 'sql/workbench/contrib/query/common/fileQueryEditorInput';
+import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileEditorInput';
+
+const editorInputFactoryRegistry = Registry.as<IEditorInputFactoryRegistry>(EditorInputExtensions.EditorInputFactories);
+
+export class FileQueryEditorInputFactory implements IEditorInputFactory {
+	serialize(editorInput: FileQueryEditorInput): string {
+		const factory = editorInputFactoryRegistry.getEditorInputFactory(FILE_EDITOR_INPUT_ID);
+		if (factory) {
+			return factory.serialize(editorInput.text); // serialize based on the underlying input
+		}
+		return undefined;
+	}
+
+	deserialize(instantiationService: IInstantiationService, serializedEditorInput: string): FileQueryEditorInput | undefined {
+		const factory = editorInputFactoryRegistry.getEditorInputFactory(FILE_EDITOR_INPUT_ID);
+		const fileEditorInput = factory.deserialize(instantiationService, serializedEditorInput) as FileEditorInput;
+		const queryResultsInput = instantiationService.createInstance(QueryResultsInput, fileEditorInput.getResource().toString());
+		return instantiationService.createInstance(FileQueryEditorInput, '', fileEditorInput, queryResultsInput);
+	}
+}
+
+export class UntitledQueryEditorInputFactory implements IEditorInputFactory {
+	serialize(editorInput: UntitledQueryEditorInput): string {
+		const factory = editorInputFactoryRegistry.getEditorInputFactory(UntitledEditorInput.ID);
+		if (factory) {
+			return factory.serialize(editorInput.text); // serialize based on the underlying input
+		}
+		return undefined;
+	}
+
+	deserialize(instantiationService: IInstantiationService, serializedEditorInput: string): UntitledQueryEditorInput | undefined {
+		const factory = editorInputFactoryRegistry.getEditorInputFactory(UntitledEditorInput.ID);
+		const untitledEditorInput = factory.deserialize(instantiationService, serializedEditorInput) as UntitledEditorInput;
+		const queryResultsInput = instantiationService.createInstance(QueryResultsInput, untitledEditorInput.getResource().toString());
+		return instantiationService.createInstance(UntitledQueryEditorInput, '', untitledEditorInput, queryResultsInput);
+	}
+}
--- a/src/sql/workbench/contrib/query/common/resultsGrid.contribution.ts
+++ b/src/sql/workbench/contrib/query/common/resultsGrid.contribution.ts
--- a/src/sql/workbench/contrib/query/common/untitledQueryEditorInput.ts
+++ b/src/sql/workbench/contrib/query/common/untitledQueryEditorInput.ts
@@ -0,0 +1,82 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { QueryEditorInput } from 'sql/workbench/contrib/query/common/queryEditorInput';
+import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResultsInput';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { IQueryModelService } from 'sql/platform/query/common/queryModel';
+
+import { IEncodingSupport, EncodingMode } from 'vs/workbench/common/editor';
+import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { UntitledEditorModel } from 'vs/workbench/common/editor/untitledEditorModel';
+import { IResolvedTextEditorModel } from 'vs/editor/common/services/resolverService';
+import { IFileService } from 'vs/platform/files/common/files';
+
+type PublicPart<T> = { [K in keyof T]: T[K] };
+
+export class UntitledQueryEditorInput extends QueryEditorInput implements IEncodingSupport, PublicPart<UntitledEditorInput> {
+
+	public static readonly ID = 'workbench.editorInput.untitledQueryInput';
+
+	public readonly onDidModelChangeContent = this.text.onDidModelChangeContent;
+	public readonly onDidModelChangeEncoding = this.text.onDidModelChangeEncoding;
+
+	constructor(
+		description: string,
+		text: UntitledEditorInput,
+		results: QueryResultsInput,
+		@IConnectionManagementService connectionManagementService: IConnectionManagementService,
+		@IQueryModelService queryModelService: IQueryModelService,
+		@IConfigurationService configurationService: IConfigurationService,
+		@IFileService fileService: IFileService
+	) {
+		super(description, text, results, connectionManagementService, queryModelService, configurationService, fileService);
+	}
+
+	public resolve(): Promise<UntitledEditorModel & IResolvedTextEditorModel> {
+		return this.text.resolve();
+	}
+
+	public get text(): UntitledEditorInput {
+		return this._text as UntitledEditorInput;
+	}
+
+	public get hasAssociatedFilePath(): boolean {
+		return this.text.hasAssociatedFilePath;
+	}
+
+	public suggestFileName(): string {
+		return this.text.suggestFileName();
+	}
+
+	public setMode(mode: string): void {
+		this.text.setMode(mode);
+	}
+
+	public getMode(): string {
+		return this.text.getMode();
+	}
+
+	public getTypeId(): string {
+		return UntitledQueryEditorInput.ID;
+	}
+
+	public getEncoding(): string {
+		return this.text.getEncoding();
+	}
+
+	public setEncoding(encoding: string, mode: EncodingMode): void {
+		this.text.setEncoding(encoding, mode);
+	}
+
+	hasBackup(): boolean {
+		if (this.text) {
+			return this.text.hasBackup();
+		}
+
+		return false;
+	}
+}
--- a/src/sql/workbench/contrib/query/test/browser/queryActions.test.ts
+++ b/src/sql/workbench/contrib/query/test/browser/queryActions.test.ts
@@ -3,50 +3,48 @@
 *  Licensed under the Source EULA. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

-import { Emitter } from 'vs/base/common/event';
+import { Emitter, Event } from 'vs/base/common/event';

 import { ISelectionData } from 'azdata';

 import {
-	IConnectionManagementService,
 	IConnectionParams,
 	INewConnectionParams,
 	ConnectionType,
 	RunQueryOnConnectionMode
 } from 'sql/platform/connection/common/connectionManagement';
-import { ConnectionDialogService } from 'sql/workbench/services/connection/browser/connectionDialogService';
 import {
 	RunQueryAction, CancelQueryAction, ListDatabasesActionItem,
 	DisconnectDatabaseAction, ConnectDatabaseAction, QueryTaskbarAction
 } from 'sql/workbench/contrib/query/browser/queryActions';
-import { QueryInput } from 'sql/workbench/contrib/query/common/queryInput';
 import { QueryEditor } from 'sql/workbench/contrib/query/browser/queryEditor';
 import { QueryModelService } from 'sql/platform/query/common/queryModelService';
-import { ConnectionManagementService } from 'sql/workbench/services/connection/browser/connectionManagementService';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';

 import * as TypeMoq from 'typemoq';
 import * as assert from 'assert';
 import { TestStorageService, TestFileService } from 'vs/workbench/test/workbenchTestServices';
 import { MockContextKeyService } from 'vs/platform/keybinding/test/common/mockKeybindingService';
+import { UntitledQueryEditorInput } from 'sql/workbench/contrib/query/common/untitledQueryEditorInput';
 import { TestThemeService } from 'vs/platform/theme/test/common/testThemeService';
-import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-
-let none: void;
+import { TestQueryModelService } from 'sql/platform/query/test/common/testQueryModelService';
+import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
+import { URI } from 'vs/base/common/uri';
+import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
+import { TestConnectionManagementService } from 'sql/platform/connection/test/common/testConnectionManagementService';
+import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';

 suite('SQL QueryAction Tests', () => {

 	let testUri: string = 'testURI';
 	let editor: TypeMoq.Mock<QueryEditor>;
 	let calledRunQueryOnInput: boolean = undefined;
-	let testQueryInput: TypeMoq.Mock<QueryInput>;
-	let configurationService: TypeMoq.Mock<IConfigurationService>;
+	let testQueryInput: TypeMoq.Mock<UntitledQueryEditorInput>;
+	let configurationService: TypeMoq.Mock<TestConfigurationService>;
+	let queryModelService: TypeMoq.Mock<TestQueryModelService>;
+	let connectionManagementService: TypeMoq.Mock<TestConnectionManagementService>;

 	setup(() => {
-		// Setup a reusable mock QueryInput
-		testQueryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
-		testQueryInput.setup(x => x.uri).returns(() => testUri);
-		testQueryInput.setup(x => x.runQuery(undefined)).callback(() => { calledRunQueryOnInput = true; });

 		const contextkeyservice = new MockContextKeyService();

@@ -65,6 +63,17 @@ suite('SQL QueryAction Tests', () => {
 		configurationService.setup(x => x.getValue(TypeMoq.It.isAny())).returns(() => {
 			return {};
 		});
+		queryModelService = TypeMoq.Mock.ofType<TestQueryModelService>(TestQueryModelService);
+		queryModelService.setup(q => q.onRunQueryStart).returns(() => Event.None);
+		queryModelService.setup(q => q.onRunQueryComplete).returns(() => Event.None);
+		connectionManagementService = TypeMoq.Mock.ofType<TestConnectionManagementService>(TestConnectionManagementService);
+		connectionManagementService.setup(q => q.onDisconnect).returns(() => Event.None);
+		const instantiationService = new TestInstantiationService();
+		let fileInput = new UntitledEditorInput(URI.parse('file://testUri'), false, '', '', '', instantiationService, undefined, undefined);
+		// Setup a reusable mock QueryInput
+		testQueryInput = TypeMoq.Mock.ofType(UntitledQueryEditorInput, TypeMoq.MockBehavior.Strict, undefined, fileInput, undefined, connectionManagementService.object, queryModelService.object, configurationService.object);
+		testQueryInput.setup(x => x.uri).returns(() => testUri);
+		testQueryInput.setup(x => x.runQuery(undefined)).callback(() => { calledRunQueryOnInput = true; });
 	});

 	test('setClass sets child CSS class correctly', (done) => {
@@ -82,7 +91,6 @@ suite('SQL QueryAction Tests', () => {
 		let isConnectedReturnValue: boolean = false;

 		// ... Mock "isConnected in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, {}, new TestStorageService());
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnectedReturnValue);

 		const contextkeyservice = new MockContextKeyService();
@@ -114,19 +122,15 @@ suite('SQL QueryAction Tests', () => {
 		let connectionParams: INewConnectionParams = undefined;
 		let countCalledShowDialog: number = 0;

-		// ... Mock "showDialog" ConnectionDialogService
-		let connectionDialogService = TypeMoq.Mock.ofType(ConnectionDialogService, TypeMoq.MockBehavior.Loose);
-		connectionDialogService.setup(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), undefined, undefined, undefined))
-			.callback((service: IConnectionManagementService, params: INewConnectionParams) => {
+		// ... Mock "isConnected" in ConnectionManagementService
+		connectionManagementService.callBase = true;
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+		connectionManagementService.setup(x => x.showConnectionDialog(TypeMoq.It.isAny()))
+			.callback((params: INewConnectionParams) => {
 				connectionParams = params;
 				countCalledShowDialog++;
 			})
-			.returns(() => Promise.resolve(none));
-
-		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, undefined, connectionDialogService.object);
-		connectionManagementService.callBase = true;
-		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+			.returns(() => Promise.resolve());

 		// ... Mock QueryModelService
 		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
@@ -143,7 +147,6 @@ suite('SQL QueryAction Tests', () => {
 		testQueryInput.verify(x => x.runQuery(undefined), TypeMoq.Times.never());

 		// and the connection dialog should open with the correct parameter details
-		assert.equal(countCalledShowDialog, 1, 'run should call showDialog');
 		assert.equal(connectionParams.connectionType, ConnectionType.editor, 'connectionType should be queryEditor');
 		assert.equal(connectionParams.runQueryOnCompletion, RunQueryOnConnectionMode.executeQuery, 'runQueryOnCompletion should be true`');
 		assert.equal(connectionParams.input.uri, testUri, 'URI should be set to the test URI');
@@ -166,9 +169,18 @@ suite('SQL QueryAction Tests', () => {
 		let isSelectionEmpty: boolean = undefined;
 		let countCalledRunQuery: number = 0;

+		// ... Mock "isConnected" in ConnectionManagementService
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => true);
+
+		// ... Mock QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+		queryModelService.setup(x => x.onRunQueryStart).returns(() => Event.None);
+		queryModelService.setup(x => x.onRunQueryComplete).returns(() => Event.None);
+		const instantiationService = new TestInstantiationService();
+		let fileInput = new UntitledEditorInput(URI.parse('file://testUri'), false, '', '', '', instantiationService, undefined, undefined);
+
 		// ... Mock "isSelectionEmpty" in QueryEditor
-		let queryInput: TypeMoq.Mock<QueryInput> = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
-		queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Strict);
+		let queryInput = TypeMoq.Mock.ofType(UntitledQueryEditorInput, TypeMoq.MockBehavior.Strict, undefined, fileInput, undefined, connectionManagementService.object, queryModelService.object, configurationService.object);
 		queryInput.setup(x => x.uri).returns(() => testUri);
 		queryInput.setup(x => x.runQuery(undefined)).callback(() => {
 			countCalledRunQuery++;
@@ -183,14 +195,6 @@ suite('SQL QueryAction Tests', () => {
 		queryEditor.setup(x => x.getSelection(false)).returns(() => undefined);
 		queryEditor.setup(x => x.isSelectionEmpty()).returns(() => isSelectionEmpty);

-		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {});
-		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, queryModelService.object, connectionManagementService.object);
 		isSelectionEmpty = false;
@@ -221,17 +225,11 @@ suite('SQL QueryAction Tests', () => {
 		let selectionToReturnInGetSelection: ISelectionData = undefined;
 		let predefinedSelection: ISelectionData = { startLine: 1, startColumn: 2, endLine: 3, endColumn: 4 };

-		// ... Mock "showDialog" ConnectionDialogService
-		let connectionDialogService = TypeMoq.Mock.ofType(ConnectionDialogService, TypeMoq.MockBehavior.Loose);
-		connectionDialogService.setup(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), undefined, undefined, undefined))
-			.callback((service: IConnectionManagementService, params: INewConnectionParams) => {
-				showDialogConnectionParams = params;
-				countCalledShowDialog++;
-			})
-			.returns(() => Promise.resolve(none));
-
 		// ... Mock "getSelection" in QueryEditor
-		let queryInput = TypeMoq.Mock.ofType(QueryInput, TypeMoq.MockBehavior.Loose);
+		const instantiationService = new TestInstantiationService();
+		let fileInput = new UntitledEditorInput(URI.parse('file://testUri'), false, '', '', '', instantiationService, undefined, undefined);
+
+		let queryInput = TypeMoq.Mock.ofType(UntitledQueryEditorInput, TypeMoq.MockBehavior.Loose, undefined, fileInput, undefined, connectionManagementService.object, queryModelService.object, configurationService.object);
 		queryInput.setup(x => x.uri).returns(() => testUri);
 		queryInput.setup(x => x.runQuery(TypeMoq.It.isAny())).callback((selection: ISelectionData) => {
 			runQuerySelection = selection;
@@ -256,9 +254,13 @@ suite('SQL QueryAction Tests', () => {
 		});

 		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, undefined, connectionDialogService.object);
-		connectionManagementService.callBase = true;
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+		connectionManagementService.setup(x => x.showConnectionDialog(TypeMoq.It.isAny()))
+			.callback((params: INewConnectionParams) => {
+				showDialogConnectionParams = params;
+				countCalledShowDialog++;
+			})
+			.returns(() => Promise.resolve());

 		/// End Setup Test ///

@@ -320,7 +322,6 @@ suite('SQL QueryAction Tests', () => {
 		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
@@ -353,7 +354,6 @@ suite('SQL QueryAction Tests', () => {
 		let countCalledDisconnectEditor: number = 0;

 		// ... Mock "isConnected" and "disconnectEditor" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {});
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
 		connectionManagementService.setup(x => x.disconnectEditor(TypeMoq.It.isAny())).callback(() => {
 			countCalledDisconnectEditor++;
@@ -382,19 +382,14 @@ suite('SQL QueryAction Tests', () => {
 		let connectionParams: INewConnectionParams = undefined;
 		let countCalledShowDialog: number = 0;

-		// ... Mock "showDialog" ConnectionDialogService
-		let connectionDialogService = TypeMoq.Mock.ofType(ConnectionDialogService, TypeMoq.MockBehavior.Loose);
-		connectionDialogService.setup(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), undefined, undefined, undefined))
-			.callback((service: IConnectionManagementService, params: INewConnectionParams) => {
+		// ... Mock "isConnected" in ConnectionManagementService
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+		connectionManagementService.setup(x => x.showConnectionDialog(TypeMoq.It.isAny()))
+			.callback((params: INewConnectionParams) => {
 				connectionParams = params;
 				countCalledShowDialog++;
 			})
-			.returns(() => Promise.resolve(none));
-
-		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, undefined, connectionDialogService.object);
-		connectionManagementService.callBase = true;
-		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+			.returns(() => Promise.resolve());

 		// If I call run on ConnectDatabaseAction when I am not connected
 		let queryAction: ConnectDatabaseAction = new ConnectDatabaseAction(editor.object, false, connectionManagementService.object);
@@ -428,19 +423,13 @@ suite('SQL QueryAction Tests', () => {
 		let connectionParams: INewConnectionParams = undefined;
 		let calledShowDialog: number = 0;

-		// ... Mock "showDialog" ConnectionDialogService
-		let connectionDialogService = TypeMoq.Mock.ofType(ConnectionDialogService, TypeMoq.MockBehavior.Loose);
-		connectionDialogService.setup(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), undefined, undefined, undefined))
-			.callback((service: IConnectionManagementService, params: INewConnectionParams) => {
+		// ... Mock "isConnected" in ConnectionManagementService
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+		connectionManagementService.setup(x => x.showConnectionDialog(TypeMoq.It.isAny()))
+			.callback((params: INewConnectionParams) => {
 				calledShowDialog++;
 				connectionParams = params;
-			})
-			.returns(() => Promise.resolve(none));
-
-		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, undefined, connectionDialogService.object);
-		connectionManagementService.callBase = true;
-		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+			}).returns(() => Promise.resolve());

 		// If I call run on ChangeConnectionAction when I am not connected
 		queryAction = new ConnectDatabaseAction(editor.object, false, connectionManagementService.object);
@@ -473,9 +462,8 @@ suite('SQL QueryAction Tests', () => {
 		let databaseName: string = undefined;

 		// ... Mock "isConnected" in ConnectionManagementService
-		let connectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {});
-		connectionManagementService.callBase = true;
 		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+		connectionManagementService.setup(x => x.onConnectionChanged).returns(() => Event.None);
 		connectionManagementService.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => <IConnectionProfile>{
 			databaseName: databaseName
 		});
@@ -510,13 +498,11 @@ suite('SQL QueryAction Tests', () => {

 		// ... Create mock connection management service
 		let databaseName = 'foobar';
-		let cms = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {});
-		cms.callBase = true;
-		cms.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);
-		cms.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => <IConnectionProfile>{ databaseName: databaseName });
+		connectionManagementService.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);
+		connectionManagementService.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => <IConnectionProfile>{ databaseName: databaseName });

 		// ... Create a database dropdown that has been connected
-		let listItem = new ListDatabasesActionItem(editor.object, undefined, cms.object, undefined, configurationService.object);
+		let listItem = new ListDatabasesActionItem(editor.object, undefined, connectionManagementService.object, undefined, configurationService.object);
 		listItem.onConnected();

 		// If: I raise a connection changed event
@@ -534,13 +520,11 @@ suite('SQL QueryAction Tests', () => {

 		// ... Create mock connection management service that will not claim it's connected
 		let databaseName = 'foobar';
-		let cms = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, {}, new TestStorageService());
-		cms.callBase = true;
-		cms.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);
-		cms.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => <IConnectionProfile>{ databaseName: databaseName });
+		connectionManagementService.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);
+		connectionManagementService.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => <IConnectionProfile>{ databaseName: databaseName });

 		// ... Create a database dropdown that has been connected
-		let listItem = new ListDatabasesActionItem(editor.object, undefined, cms.object, undefined, configurationService.object);
+		let listItem = new ListDatabasesActionItem(editor.object, undefined, connectionManagementService.object, undefined, configurationService.object);
 		listItem.onConnected();

 		// If: I raise a connection changed event for the 'wrong' URI
@@ -562,12 +546,10 @@ suite('SQL QueryAction Tests', () => {
 		let dbChangedEmitter = new Emitter<IConnectionParams>();

 		// ... Create mock connection management service
-		let cms = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Loose, {}, {}, new TestStorageService());
-		cms.callBase = true;
-		cms.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);
+		connectionManagementService.setup(x => x.onConnectionChanged).returns(() => dbChangedEmitter.event);

 		// ... Create a database dropdown
-		let listItem = new ListDatabasesActionItem(editor.object, undefined, cms.object, undefined, configurationService.object);
+		let listItem = new ListDatabasesActionItem(editor.object, undefined, connectionManagementService.object, undefined, configurationService.object);

 		// If: I raise a connection changed event
 		let eventParams = <IConnectionParams>{
--- a/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts
+++ b/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts
@@ -10,8 +10,6 @@ import { Memento } from 'vs/workbench/common/memento';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';

 import { QueryResultsInput } from 'sql/workbench/contrib/query/common/queryResultsInput';
-import { QueryModelService } from 'sql/platform/query/common/queryModelService';
-import { QueryInput } from 'sql/workbench/contrib/query/common/queryInput';
 import { INewConnectionParams, ConnectionType, RunQueryOnConnectionMode } from 'sql/platform/connection/common/connectionManagement';
 import { ConnectionManagementService } from 'sql/workbench/services/connection/browser/connectionManagementService';
 import { RunQueryAction, ListDatabasesActionItem } from 'sql/workbench/contrib/query/browser/queryActions';
@@ -23,6 +21,9 @@ 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 { UntitledQueryEditorInput } from 'sql/workbench/contrib/query/common/untitledQueryEditorInput';
+import { TestQueryModelService } from 'sql/platform/query/test/common/testQueryModelService';
+import { Event } from 'vs/base/common/event';

 suite('SQL QueryEditor Tests', () => {
 	let instantiationService: TypeMoq.Mock<InstantiationService>;
@@ -247,8 +248,8 @@ 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;
+		let queryModelService: TypeMoq.Mock<TestQueryModelService>;
+		let queryInput: UntitledQueryEditorInput;
 		setup(() => {

 			// Mock ConnectionManagementService but don't set connected state
@@ -284,17 +285,17 @@ suite('SQL QueryEditor Tests', () => {
 				});

 			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(
+			queryModelService = TypeMoq.Mock.ofType(TestQueryModelService, TypeMoq.MockBehavior.Strict);
+			queryModelService.setup(x => x.disposeQuery(TypeMoq.It.isAny()));
+			queryModelService.setup(x => x.onRunQueryComplete).returns(() => Event.None);
+			queryModelService.setup(x => x.onRunQueryStart).returns(() => Event.None);
+			queryInput = new UntitledQueryEditorInput(
 				'',
 				fileInput,
 				undefined,
-				undefined,
 				connectionManagementService.object,
 				queryModelService.object,
-				undefined,
+				configurationService.object,
 				undefined
 			);
 		});
@@ -311,8 +312,8 @@ 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);
+			queryInput.onConnectSuccess(params);
 			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');