Strict null on some query and connection (#7300)

* wip

* make connection work with strict-nulls

* change comments

* fix tests; remove unneeded type forcing

* address feedback

* adjust the logic of query editor

* clean up typing

src/sql/platform/common/pathUtilities.ts

@@ -11,7 +11,7 @@ import { Schemas } from 'vs/base/common/network';
 
 export const FILE_SCHEMA: string = 'file';
 
-export function resolveCurrentDirectory(uri: string, rootPath: string): string | undefined {
+export function resolveCurrentDirectory(uri: string, rootPath?: string): string | undefined {
 	let sqlUri = URI.parse(uri);
 	let currentDirectory: string | undefined;
 
@@ -30,7 +30,7 @@ export function resolveCurrentDirectory(uri: string, rootPath: string): string |
 	return currentDirectory;
 }
 
-export function resolveFilePath(uri: string, filePath: string, rootPath: string): string | undefined {
+export function resolveFilePath(uri: string, filePath: string, rootPath?: string): string | undefined {
 	let currentDirectory = resolveCurrentDirectory(uri, rootPath);
 	if (currentDirectory) {
 		return normalize(join(currentDirectory, filePath));

src/sql/platform/connection/common/connectionConfig.ts

@@ -18,7 +18,7 @@ const CONNECTIONS_CONFIG_KEY = 'datasource.connections';
 
 export interface ISaveGroupResult {
 	groups: IConnectionProfileGroup[];
-	newGroupId: string;
+	newGroupId?: string;
 }
 
 /**
@@ -37,7 +37,9 @@ export class ConnectionConfig {
 	public getAllGroups(): IConnectionProfileGroup[] {
 
 		let allGroups: IConnectionProfileGroup[] = [];
-		let { user, workspace } = this.configurationService.inspect<IConnectionProfileGroup[]>(GROUPS_CONFIG_KEY);
+		const config = this.configurationService.inspect<IConnectionProfileGroup[]>(GROUPS_CONFIG_KEY);
+		let { user } = config;
+		const { workspace } = config;
 
 		if (user) {
 			if (workspace) {
@@ -109,7 +111,7 @@ export class ConnectionConfig {
 			let result = this.saveGroup(groups, profile.groupFullName, undefined, undefined);
 			groups = result.groups;
 
-			return this.configurationService.updateValue(GROUPS_CONFIG_KEY, groups, ConfigurationTarget.USER).then(() => result.newGroupId);
+			return this.configurationService.updateValue(GROUPS_CONFIG_KEY, groups, ConfigurationTarget.USER).then(() => result.newGroupId!);
 		}
 	}
 
@@ -129,7 +131,7 @@ export class ConnectionConfig {
 				let result = this.saveGroup(groups, profileGroup.name, profileGroup.color, profileGroup.description);
 				groups = result.groups;
 
-				return this.configurationService.updateValue(GROUPS_CONFIG_KEY, groups, ConfigurationTarget.USER).then(() => result.newGroupId);
+				return this.configurationService.updateValue(GROUPS_CONFIG_KEY, groups, ConfigurationTarget.USER).then(() => result.newGroupId!);
 			}
 		}
 	}
@@ -138,18 +140,22 @@ export class ConnectionConfig {
 		let configs = this.configurationService.inspect<IConnectionProfileStore[]>(CONNECTIONS_CONFIG_KEY);
 		let profiles: IConnectionProfileStore[];
 		if (configs) {
+			let fromConfig: IConnectionProfileStore[] | undefined;
 			if (configTarget === ConfigurationTarget.USER) {
-				profiles = configs.user;
+				fromConfig = configs.user;
 			} else if (configTarget === ConfigurationTarget.WORKSPACE) {
-				profiles = configs.workspace;
+				fromConfig = configs.workspace || [];
 			}
-			if (profiles) {
+			if (fromConfig) {
+				profiles = fromConfig;
 				if (this.fixConnectionIds(profiles)) {
 					this.configurationService.updateValue(CONNECTIONS_CONFIG_KEY, profiles, configTarget);
 				}
 			} else {
 				profiles = [];
 			}
+		} else {
+			profiles = [];
 		}
 
 		return profiles;
@@ -315,7 +321,7 @@ export class ConnectionConfig {
 		}
 	}
 
-	public saveGroup(groups: IConnectionProfileGroup[], groupFullName: string, color: string, description: string): ISaveGroupResult {
+	public saveGroup(groups: IConnectionProfileGroup[], groupFullName?: string, color?: string, description?: string): ISaveGroupResult {
 		let groupNames = ConnectionProfileGroup.getGroupFullNameParts(groupFullName);
 		return this.saveGroupInTree(groups, undefined, groupNames, color, description, 0);
 	}
@@ -348,21 +354,21 @@ export class ConnectionConfig {
 		return sameGroupName;
 	}
 
-	private saveGroupInTree(groupTree: IConnectionProfileGroup[], parentId: string, groupNames: string[], color: string, description: string, index: number): ISaveGroupResult {
+	private saveGroupInTree(groupTree: IConnectionProfileGroup[], parentId: string | undefined, groupNames: string[], color: string | undefined, description: string | undefined, index: number): ISaveGroupResult {
 		if (!groupTree) {
 			groupTree = [];
 		}
-		let newGroupId: string;
+		let newGroupId: string | undefined;
 
 		if (index < groupNames.length) {
 			let groupName: string = groupNames[index];
-			let newGroup: IConnectionProfileGroup = {
+			let newGroup = <unknown>{ // workaround to make this work properly
 				name: groupName,
 				id: undefined,
 				parentId: parentId,
 				color: color,
 				description: description
-			};
+			} as IConnectionProfileGroup;
 			let found = groupTree.find(group => this.isSameGroupName(group, newGroup));
 			if (found) {
 				if (index === groupNames.length - 1) {

src/sql/platform/connection/common/connectionProfile.ts

@@ -25,7 +25,7 @@ export class ConnectionProfile extends ProviderConnectionInfo implements interfa
 	public parent?: ConnectionProfileGroup;
 	private _id: string;
 	public savePassword: boolean;
-	private _groupName: string;
+	private _groupName?: string;
 	public groupId: string;
 	public saveProfile: boolean;
 
@@ -119,11 +119,11 @@ export class ConnectionProfile extends ProviderConnectionInfo implements interfa
 		this.options['registeredServerDescription'] = value;
 	}
 
-	public get groupFullName(): string {
+	public get groupFullName(): string | undefined {
 		return this._groupName;
 	}
 
-	public set groupFullName(value: string) {
+	public set groupFullName(value: string | undefined) {
 		this._groupName = value;
 	}
 
@@ -214,15 +214,12 @@ export class ConnectionProfile extends ProviderConnectionInfo implements interfa
 		return !profile.databaseName || profile.databaseName.trim() === '';
 	}
 
-	public static fromIConnectionProfile(capabilitiesService: ICapabilitiesService, profile: azdata.IConnectionProfile) {
-		if (profile) {
-			if (profile instanceof ConnectionProfile) {
-				return profile;
-			} else {
-				return new ConnectionProfile(capabilitiesService, profile);
-			}
+	public static fromIConnectionProfile(capabilitiesService: ICapabilitiesService, profile: azdata.IConnectionProfile): ConnectionProfile {
+		if (profile instanceof ConnectionProfile) {
+			return profile;
+		} else {
+			return new ConnectionProfile(capabilitiesService, profile);
 		}
-		return undefined;
 	}
 
 	public static createFromStoredProfile(profile: interfaces.IConnectionProfileStore, capabilitiesService: ICapabilitiesService): ConnectionProfile {
@@ -244,24 +241,18 @@ export class ConnectionProfile extends ProviderConnectionInfo implements interfa
 
 	public static convertToProfileStore(
 		capabilitiesService: ICapabilitiesService,
-		connectionProfile: interfaces.IConnectionProfile): interfaces.IConnectionProfileStore | undefined {
-		if (connectionProfile) {
-			let connectionInfo = ConnectionProfile.fromIConnectionProfile(capabilitiesService, connectionProfile);
-			if (connectionInfo) {
-				let profile: interfaces.IConnectionProfileStore = {
-					options: {},
-					groupId: connectionProfile.groupId,
-					providerName: connectionInfo.providerName,
-					savePassword: connectionInfo.savePassword,
-					id: connectionInfo.id
-				};
+		connectionProfile: interfaces.IConnectionProfile): interfaces.IConnectionProfileStore {
+		let connectionInfo = ConnectionProfile.fromIConnectionProfile(capabilitiesService, connectionProfile);
+		let profile: interfaces.IConnectionProfileStore = {
+			options: {},
+			groupId: connectionProfile.groupId,
+			providerName: connectionInfo.providerName,
+			savePassword: connectionInfo.savePassword,
+			id: connectionInfo.id
+		};
 
-				profile.options = connectionInfo.options;
+		profile.options = connectionInfo.options;
 
-				return profile;
-			}
-		}
-
-		return undefined;
+		return profile;
 	}
 }

src/sql/platform/connection/common/connectionProfileGroup.ts

@@ -5,13 +5,14 @@
 
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { Disposable } from 'vs/base/common/lifecycle';
+import { isUndefinedOrNull } from 'vs/base/common/types';
 
 export interface IConnectionProfileGroup {
 	id: string;
 	parentId?: string;
 	name: string;
-	color: string;
-	description: string;
+	color?: string;
+	description?: string;
 }
 
 export class ConnectionProfileGroup extends Disposable implements IConnectionProfileGroup {
@@ -22,10 +23,10 @@ export class ConnectionProfileGroup extends Disposable implements IConnectionPro
 	private _isRenamed: boolean;
 	public constructor(
 		public name: string,
-		public parent: ConnectionProfileGroup,
+		public parent: ConnectionProfileGroup | undefined,
 		public id: string,
-		public color: string,
-		public description: string
+		public color?: string,
+		public description?: string
 	) {
 		super();
 		this.parentId = parent ? parent.id : undefined;
@@ -150,7 +151,7 @@ export class ConnectionProfileGroup extends Disposable implements IConnectionPro
 		});
 	}
 
-	public getParent(): ConnectionProfileGroup {
+	public getParent(): ConnectionProfileGroup | undefined {
 		return this.parent;
 	}
 
@@ -167,7 +168,7 @@ export class ConnectionProfileGroup extends Disposable implements IConnectionPro
 		return isAncestor;
 	}
 
-	public static getGroupFullNameParts(groupFullName: string): string[] {
+	public static getGroupFullNameParts(groupFullName?: string): string[] {
 		groupFullName = groupFullName ? groupFullName : '';
 		let groupNames: string[] = groupFullName.split(ConnectionProfileGroup.GroupNameSeparator);
 		groupNames = groupNames.filter(g => !!g);
@@ -185,13 +186,19 @@ export class ConnectionProfileGroup extends Disposable implements IConnectionPro
 			name === ConnectionProfileGroup.GroupNameSeparator);
 	}
 
-	public static sameGroupName(name1: string, name2: string): boolean {
-		let sameGroupName: boolean =
-			(!name1 && !name2) ||
-			name1.toUpperCase() === name2.toUpperCase() ||
-			(ConnectionProfileGroup.isRoot(name1) && ConnectionProfileGroup.isRoot(name2));
-
-		return sameGroupName;
+	public static sameGroupName(name1?: string, name2?: string): boolean {
+		const isName1Undefined = isUndefinedOrNull(name1);
+		const isName2Undefined = isUndefinedOrNull(name2);
+		if (isName1Undefined && isName2Undefined) {
+			return true;
+		}
+		if ((isName1Undefined && !isName2Undefined) || !isName1Undefined && isName2Undefined) {
+			return false;
+		}
+		if (name1!.toUpperCase() === name2!.toUpperCase()) {
+			return true;
+		}
+		return ConnectionProfileGroup.isRoot(name1!) && ConnectionProfileGroup.isRoot(name2!);
 	}
 
 	public static getConnectionsInGroup(group: ConnectionProfileGroup): ConnectionProfile[] {

src/sql/platform/connection/common/connectionStatusManager.ts

@@ -27,7 +27,7 @@ export class ConnectionStatusManager {
 		this._connections = {};
 	}
 
-	public findConnection(uri: string): ConnectionManagementInfo {
+	public findConnection(uri: string): ConnectionManagementInfo | undefined {
 		if (uri in this._connections) {
 			return this._connections[uri];
 		} else {
@@ -39,7 +39,7 @@ export class ConnectionStatusManager {
 		return Object.values(this._connections).find((connection: ConnectionManagementInfo) => connection.connectionProfile.id === profileId);
 	}
 
-	public findConnectionProfile(connectionProfile: IConnectionProfile): ConnectionManagementInfo {
+	public findConnectionProfile(connectionProfile: IConnectionProfile): ConnectionManagementInfo | undefined {
 		let id = Utils.generateUri(connectionProfile);
 		return this.findConnection(id);
 	}
@@ -65,7 +65,7 @@ export class ConnectionStatusManager {
 		}
 	}
 
-	public getConnectionProfile(id: string): ConnectionProfile {
+	public getConnectionProfile(id: string): ConnectionProfile | undefined {
 		let connectionInfoForId = this.findConnection(id);
 		return connectionInfoForId ? connectionInfoForId.connectionProfile : undefined;
 	}
@@ -178,7 +178,7 @@ export class ConnectionStatusManager {
 		return ownerUriToReturn;
 	}
 
-	public onConnectionChanged(changedConnInfo: azdata.ChangedConnectionInfo): IConnectionProfile {
+	public onConnectionChanged(changedConnInfo: azdata.ChangedConnectionInfo): IConnectionProfile | undefined {
 		let connection = this._connections[changedConnInfo.connectionUri];
 		if (connection && connection.connectionProfile) {
 			connection.connectionProfile.serverName = changedConnInfo.connection.serverName;
@@ -190,14 +190,14 @@ export class ConnectionStatusManager {
 	}
 
 	private isSharedSession(fileUri: string): boolean {
-		return fileUri && fileUri.startsWith('vsls:');
+		return !!(fileUri && fileUri.startsWith('vsls:'));
 	}
 
 	public isConnected(id: string): boolean {
 		if (this.isSharedSession(id)) {
 			return true;
 		}
-		return (id in this._connections && this._connections[id].connectionId && !!this._connections[id].connectionId);
+		return !!(id in this._connections && this._connections[id].connectionId && !!this._connections[id].connectionId);
 	}
 
 	public isConnecting(id: string): boolean {
@@ -205,7 +205,7 @@ export class ConnectionStatusManager {
 	}
 
 	public isDefaultTypeUri(uri: string): boolean {
-		return uri && uri.startsWith(Utils.uriPrefixes.default);
+		return !!(uri && uri.startsWith(Utils.uriPrefixes.default));
 	}
 
 	public getProviderIdFromUri(ownerUri: string): string {

src/sql/platform/connection/common/connectionStore.ts

@@ -29,7 +29,7 @@ const CRED_PROFILE_USER = 'Profile';
  * @export
  */
 export class ConnectionStore {
-	private groupIdMap = new ReverseLookUpMap<string, string>();
+	private groupIdMap = new ReverseLookUpMap<string, string | undefined>();
 	private connectionConfig = new ConnectionConfig(this.configurationService, this.capabilitiesService);
 	private mru: Array<IConnectionProfile>;
 
@@ -57,8 +57,7 @@ export class ConnectionStore {
 	 * @returns formatted string with server, DB and username
 	 */
 	private formatCredentialId(connectionProfile: IConnectionProfile, itemType?: string): string {
-		const connectionProfileInstance: ConnectionProfile = ConnectionProfile.fromIConnectionProfile(
-			this.capabilitiesService, connectionProfile);
+		const connectionProfileInstance = ConnectionProfile.fromIConnectionProfile(this.capabilitiesService, connectionProfile);
 		const cred: string[] = [CRED_PREFIX];
 		if (!itemType) {
 			itemType = CRED_PROFILE_USER;
@@ -74,12 +73,8 @@ export class ConnectionStore {
 	 * @param connection profile
 	 */
 	public isPasswordRequired(connection: IConnectionProfile): boolean {
-		if (connection) {
-			const connectionProfile = ConnectionProfile.fromIConnectionProfile(this.capabilitiesService, connection);
-			return connectionProfile.isPasswordRequired();
-		} else {
-			return false;
-		}
+		const connectionProfile = ConnectionProfile.fromIConnectionProfile(this.capabilitiesService, connection);
+		return connectionProfile.isPasswordRequired();
 	}
 
 	public addSavedPassword(credentialsItem: IConnectionProfile): Promise<{ profile: IConnectionProfile, savedCred: boolean }> {
@@ -162,34 +157,26 @@ export class ConnectionStore {
 
 	private convertConfigValuesToConnectionProfiles(configValues: IConnectionProfile[]): ConnectionProfile[] {
 		return configValues.map(c => {
-			if (c) {
-				const connectionProfile = new ConnectionProfile(this.capabilitiesService, c);
-				if (connectionProfile.saveProfile) {
-					if (!connectionProfile.groupFullName && connectionProfile.groupId) {
-						connectionProfile.groupFullName = this.getGroupFullName(connectionProfile.groupId);
-					}
-					if (!connectionProfile.groupId && connectionProfile.groupFullName) {
-						connectionProfile.groupId = this.getGroupId(connectionProfile.groupFullName);
-					} else if (!connectionProfile.groupId && !connectionProfile.groupFullName) {
-						connectionProfile.groupId = this.getGroupId('');
-					}
+			const connectionProfile = new ConnectionProfile(this.capabilitiesService, c);
+			if (connectionProfile.saveProfile) {
+				if (!connectionProfile.groupFullName && connectionProfile.groupId) {
+					connectionProfile.groupFullName = this.getGroupFullName(connectionProfile.groupId);
+				}
+				if (!connectionProfile.groupId && connectionProfile.groupFullName) {
+					connectionProfile.groupId = this.getGroupId(connectionProfile.groupFullName);
+				} else if (!connectionProfile.groupId && !connectionProfile.groupFullName) {
+					connectionProfile.groupId = this.getGroupId('');
 				}
-				return connectionProfile;
-			} else {
-				return undefined;
 			}
+			return connectionProfile;
 		});
 	}
 
 	public getProfileWithoutPassword(conn: IConnectionProfile): ConnectionProfile {
-		if (conn) {
-			let savedConn: ConnectionProfile = ConnectionProfile.fromIConnectionProfile(this.capabilitiesService, conn);
-			savedConn = savedConn.withoutPassword();
+		let savedConn = ConnectionProfile.fromIConnectionProfile(this.capabilitiesService, conn);
+		savedConn = savedConn.withoutPassword();
 
-			return savedConn;
-		} else {
-			return undefined;
-		}
+		return savedConn;
 	}
 
 	/**
@@ -221,7 +208,7 @@ export class ConnectionStore {
 	}
 
 	private addToConnectionList(conn: IConnectionProfile, list: ConnectionProfile[]): IConnectionProfile[] {
-		const savedProfile: ConnectionProfile = this.getProfileWithoutPassword(conn);
+		const savedProfile = this.getProfileWithoutPassword(conn);
 
 		// Remove the connection from the list if it already exists
 		list = list.filter(value => {
@@ -239,7 +226,7 @@ export class ConnectionStore {
 	}
 
 	private removeFromConnectionList(conn: IConnectionProfile, list: ConnectionProfile[]): IConnectionProfile[] {
-		const savedProfile: ConnectionProfile = this.getProfileWithoutPassword(conn);
+		const savedProfile = this.getProfileWithoutPassword(conn);
 
 		// Remove the connection from the list if it already exists
 		list = list.filter(value => {
@@ -286,7 +273,7 @@ export class ConnectionStore {
 	}
 
 	public getConnectionProfileGroups(withoutConnections?: boolean, providers?: string[]): ConnectionProfileGroup[] {
-		let profilesInConfiguration: ConnectionProfile[];
+		let profilesInConfiguration: ConnectionProfile[] | undefined;
 		if (!withoutConnections) {
 			profilesInConfiguration = this.connectionConfig.getConnections(true);
 			if (providers && providers.length > 0) {
@@ -295,10 +282,10 @@ export class ConnectionStore {
 		}
 		const groups = this.connectionConfig.getAllGroups();
 
-		return this.convertToConnectionGroup(groups, profilesInConfiguration, undefined);
+		return this.convertToConnectionGroup(groups, profilesInConfiguration);
 	}
 
-	private convertToConnectionGroup(groups: IConnectionProfileGroup[], connections: ConnectionProfile[], parent: ConnectionProfileGroup = undefined): ConnectionProfileGroup[] {
+	private convertToConnectionGroup(groups: IConnectionProfileGroup[], connections?: ConnectionProfile[], parent?: ConnectionProfileGroup): ConnectionProfileGroup[] {
 		const result: ConnectionProfileGroup[] = [];
 		const children = groups.filter(g => g.parentId === (parent ? parent.id : undefined));
 		if (children) {
@@ -307,7 +294,7 @@ export class ConnectionStore {
 				this.addGroupFullNameToMap(group.id, connectionGroup.fullName);
 				if (connections) {
 					let connectionsForGroup = connections.filter(conn => conn.groupId === connectionGroup.id);
-					let conns = [];
+					let conns: ConnectionProfile[] = [];
 					connectionsForGroup.forEach((conn) => {
 						conn.groupFullName = connectionGroup.fullName;
 						conns.push(conn);
@@ -359,7 +346,7 @@ export class ConnectionStore {
 		return this.connectionConfig.changeGroupIdForConnection(source, targetGroupId).then();
 	}
 
-	private addGroupFullNameToMap(groupId: string, groupFullName: string): void {
+	private addGroupFullNameToMap(groupId: string, groupFullName?: string): void {
 		if (groupId) {
 			this.groupIdMap.set(groupId, groupFullName);
 		}
@@ -373,7 +360,7 @@ export class ConnectionStore {
 			// Load the cache
 			this.getConnectionProfileGroups(true);
 		}
-		return this.groupIdMap.get(groupId);
+		return this.groupIdMap.get(groupId)!;
 	}
 
 	private getGroupId(groupFullName: string): string {
@@ -385,6 +372,6 @@ export class ConnectionStore {
 			// Load the cache
 			this.getConnectionProfileGroups(true);
 		}
-		return this.groupIdMap.reverseGet(key);
+		return this.groupIdMap.reverseGet(key)!;
 	}
 }

src/sql/platform/connection/common/utils.ts

@@ -115,7 +115,7 @@ export function generateUriWithPrefix(connection: IConnectionProfile, prefix: st
 	return uri;
 }
 
-export function findProfileInGroup(og: IConnectionProfile, groups: ConnectionProfileGroup[]): ConnectionProfile {
+export function findProfileInGroup(og: IConnectionProfile, groups: ConnectionProfileGroup[]): ConnectionProfile | undefined {
 	for (let group of groups) {
 		for (let conn of group.connections) {
 			if (conn.id === og.id) {

src/sql/platform/query/common/gridDataProvider.ts

@@ -34,7 +34,7 @@ export interface IGridDataProvider {
 
 	shouldRemoveNewLines(): boolean;
 
-	getColumnHeaders(range: Slick.Range): string[];
+	getColumnHeaders(range: Slick.Range): string[] | undefined;
 
 	readonly canSerialize: boolean;
 

src/sql/platform/query/common/queryManagement.ts

@@ -29,16 +29,16 @@ export interface IQueryManagementService {
 	getRegisteredProviders(): string[];
 	registerRunner(runner: QueryRunner, uri: string): void;
 
-	cancelQuery(ownerUri: string): Thenable<azdata.QueryCancelResult>;
-	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
-	runQueryStatement(ownerUri: string, line: number, column: number): Thenable<void>;
-	runQueryString(ownerUri: string, queryString: string): Thenable<void>;
-	runQueryAndReturn(ownerUri: string, queryString: string): Thenable<azdata.SimpleExecuteResult>;
-	parseSyntax(ownerUri: string, query: string): Thenable<azdata.SyntaxParseResult>;
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Thenable<azdata.QueryExecuteSubsetResult>;
-	disposeQuery(ownerUri: string): Thenable<void>;
-	saveResults(requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult>;
-	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Thenable<void>;
+	cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult>;
+	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
+	runQueryString(ownerUri: string, queryString: string): Promise<void>;
+	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
+	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult>;
+	disposeQuery(ownerUri: string): Promise<void>;
+	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
+	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Promise<void>;
 
 	// Callbacks
 	onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void;
@@ -52,42 +52,42 @@ export interface IQueryManagementService {
 	onEditSessionReady(ownerUri: string, success: boolean, message: string): void;
 
 	// Edit Data Functions
-	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Thenable<void>;
-	disposeEdit(ownerUri: string): Thenable<void>;
-	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<azdata.EditUpdateCellResult>;
-	commitEdit(ownerUri): Thenable<void>;
-	createRow(ownerUri: string): Thenable<azdata.EditCreateRowResult>;
-	deleteRow(ownerUri: string, rowId: number): Thenable<void>;
-	revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<azdata.EditRevertCellResult>;
-	revertRow(ownerUri: string, rowId: number): Thenable<void>;
-	getEditRows(rowData: azdata.EditSubsetParams): Thenable<azdata.EditSubsetResult>;
+	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void>;
+	disposeEdit(ownerUri: string): Promise<void>;
+	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult>;
+	commitEdit(ownerUri: string): Promise<void>;
+	createRow(ownerUri: string): Promise<azdata.EditCreateRowResult>;
+	deleteRow(ownerUri: string, rowId: number): Promise<void>;
+	revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult>;
+	revertRow(ownerUri: string, rowId: number): Promise<void>;
+	getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult>;
 }
 
 /*
  * An object that can handle basic request-response actions related to queries
  */
 export interface IQueryRequestHandler {
-	cancelQuery(ownerUri: string): Thenable<azdata.QueryCancelResult>;
-	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
-	runQueryStatement(ownerUri: string, line: number, column: number): Thenable<void>;
-	runQueryString(ownerUri: string, queryString: string): Thenable<void>;
-	runQueryAndReturn(ownerUri: string, queryString: string): Thenable<azdata.SimpleExecuteResult>;
-	parseSyntax(ownerUri: string, query: string): Thenable<azdata.SyntaxParseResult>;
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Thenable<azdata.QueryExecuteSubsetResult>;
-	disposeQuery(ownerUri: string): Thenable<void>;
-	saveResults(requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult>;
-	setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Thenable<void>;
+	cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult>;
+	runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	runQueryStatement(ownerUri: string, line: number, column: number): Promise<void>;
+	runQueryString(ownerUri: string, queryString: string): Promise<void>;
+	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
+	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult>;
+	disposeQuery(ownerUri: string): Promise<void>;
+	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
+	setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void>;
 
 	// Edit Data actions
-	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Thenable<void>;
-	disposeEdit(ownerUri: string): Thenable<void>;
-	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<azdata.EditUpdateCellResult>;
-	commitEdit(ownerUri): Thenable<void>;
-	createRow(ownerUri: string): Thenable<azdata.EditCreateRowResult>;
-	deleteRow(ownerUri: string, rowId: number): Thenable<void>;
-	revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<azdata.EditRevertCellResult>;
-	revertRow(ownerUri: string, rowId: number): Thenable<void>;
-	getEditRows(rowData: azdata.EditSubsetParams): Thenable<azdata.EditSubsetResult>;
+	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void>;
+	disposeEdit(ownerUri: string): Promise<void>;
+	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult>;
+	commitEdit(ownerUri: string): Promise<void>;
+	createRow(ownerUri: string): Promise<azdata.EditCreateRowResult>;
+	deleteRow(ownerUri: string, rowId: number): Promise<void>;
+	revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult>;
+	revertRow(ownerUri: string, rowId: number): Promise<void>;
+	getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult>;
 }
 
 export class QueryManagementService implements IQueryManagementService {
@@ -117,7 +117,7 @@ export class QueryManagementService implements IQueryManagementService {
 		// _handlerCallbackQueue will be non-empty. Run all handlers in the queue first
 		// so that notifications are handled in order they arrived
 		while (this._handlerCallbackQueue.length > 0) {
-			let handler = this._handlerCallbackQueue.shift();
+			let handler = this._handlerCallbackQueue.shift()!;
 			handler(runner);
 		}
 
@@ -141,7 +141,7 @@ export class QueryManagementService implements IQueryManagementService {
 
 	private _notify(ownerUri: string, sendNotification: (runner: QueryRunner) => void): void {
 		let runner = this._queryRunners.get(ownerUri);
-		this.enqueueOrRun(sendNotification, runner);
+		this.enqueueOrRun(sendNotification, runner!);
 	}
 
 	public addQueryRequestHandler(queryType: string, handler: IQueryRequestHandler): IDisposable {
@@ -181,7 +181,7 @@ export class QueryManagementService implements IQueryManagementService {
 		TelemetryUtils.addTelemetry(this._telemetryService, this.logService, eventName, data);
 	}
 
-	private _runAction<T>(uri: string, action: (handler: IQueryRequestHandler) => Thenable<T>, fallBackToDefaultProvider: boolean = false): Thenable<T> {
+	private _runAction<T>(uri: string, action: (handler: IQueryRequestHandler) => Promise<T>, fallBackToDefaultProvider: boolean = false): Promise<T> {
 		let providerId: string = this._connectionService.getProviderIdFromUri(uri);
 
 		if (!providerId && fallBackToDefaultProvider) {
@@ -199,58 +199,58 @@ export class QueryManagementService implements IQueryManagementService {
 		}
 	}
 
-	public cancelQuery(ownerUri: string): Thenable<azdata.QueryCancelResult> {
+	public cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult> {
 		this.addTelemetry(TelemetryKeys.CancelQuery, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.cancelQuery(ownerUri);
 		});
 	}
-	public runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Thenable<void> {
+	public runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQuery, ownerUri, runOptions);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQuery(ownerUri, selection, runOptions);
 		});
 	}
-	public runQueryStatement(ownerUri: string, line: number, column: number): Thenable<void> {
+	public runQueryStatement(ownerUri: string, line: number, column: number): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQueryStatement, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryStatement(ownerUri, line, column);
 		});
 	}
-	public runQueryString(ownerUri: string, queryString: string): Thenable<void> {
+	public runQueryString(ownerUri: string, queryString: string): Promise<void> {
 		this.addTelemetry(TelemetryKeys.RunQueryString, ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryString(ownerUri, queryString);
 		});
 	}
-	public runQueryAndReturn(ownerUri: string, queryString: string): Thenable<azdata.SimpleExecuteResult> {
+	public runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryAndReturn(ownerUri, queryString);
 		});
 	}
-	public parseSyntax(ownerUri: string, query: string): Thenable<azdata.SyntaxParseResult> {
+	public parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.parseSyntax(ownerUri, query);
 		});
 	}
-	public getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Thenable<azdata.QueryExecuteSubsetResult> {
+	public getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<azdata.QueryExecuteSubsetResult> {
 		return this._runAction(rowData.ownerUri, (runner) => {
 			return runner.getQueryRows(rowData);
 		});
 	}
-	public disposeQuery(ownerUri: string): Thenable<void> {
+	public disposeQuery(ownerUri: string): Promise<void> {
 		this._queryRunners.delete(ownerUri);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.disposeQuery(ownerUri);
 		});
 	}
-	public setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Thenable<void> {
+	public setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.setQueryExecutionOptions(ownerUri, options);
 		}, true);
 	}
 
-	public saveResults(requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult> {
+	public saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult> {
 		return this._runAction(requestParams.ownerUri, (runner) => {
 			return runner.saveResults(requestParams);
 		});
@@ -292,7 +292,7 @@ export class QueryManagementService implements IQueryManagementService {
 	}
 
 	// Edit Data Functions
-	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Thenable<void> {
+	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.initializeEdit(ownerUri, schemaName, objectName, objectType, rowLimit, queryString);
 		});
@@ -304,49 +304,49 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 
-	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<azdata.EditUpdateCellResult> {
+	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.updateCell(ownerUri, rowId, columnId, newValue);
 		});
 	}
 
-	public commitEdit(ownerUri: string): Thenable<void> {
+	public commitEdit(ownerUri: string): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.commitEdit(ownerUri);
 		});
 	}
 
-	public createRow(ownerUri: string): Thenable<azdata.EditCreateRowResult> {
+	public createRow(ownerUri: string): Promise<azdata.EditCreateRowResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.createRow(ownerUri);
 		});
 	}
 
-	public deleteRow(ownerUri: string, rowId: number): Thenable<void> {
+	public deleteRow(ownerUri: string, rowId: number): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.deleteRow(ownerUri, rowId);
 		});
 	}
 
-	public disposeEdit(ownerUri: string): Thenable<void> {
+	public disposeEdit(ownerUri: string): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.disposeEdit(ownerUri);
 		});
 	}
 
-	public revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<azdata.EditRevertCellResult> {
+	public revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.revertCell(ownerUri, rowId, columnId);
 		});
 	}
 
-	public revertRow(ownerUri: string, rowId: number): Thenable<void> {
+	public revertRow(ownerUri: string, rowId: number): Promise<void> {
 		return this._runAction(ownerUri, (runner) => {
 			return runner.revertRow(ownerUri, rowId);
 		});
 	}
 
-	public getEditRows(rowData: azdata.EditSubsetParams): Thenable<azdata.EditSubsetResult> {
+	public getEditRows(rowData: azdata.EditSubsetParams): Promise<azdata.EditSubsetResult> {
 		return this._runAction(rowData.ownerUri, (runner) => {
 			return runner.getEditRows(rowData);
 		});

src/sql/platform/query/common/queryModel.ts

@@ -49,14 +49,12 @@ export interface IQueryEvent {
 export interface IQueryModelService {
 	_serviceBrand: undefined;
 
-	getQueryRunner(uri: string): QueryRunner;
+	getQueryRunner(uri: string): QueryRunner | undefined;
 
-	getConfig(): Promise<{ [key: string]: any }>;
-	getShortcuts(): Promise<any>;
-	getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Thenable<ResultSetSubset>;
-	runQuery(uri: string, selection: ISelectionData, queryInput: QueryInput, runOptions?: ExecutionPlanOptions): void;
-	runQueryStatement(uri: string, selection: ISelectionData, queryInput: QueryInput): void;
-	runQueryString(uri: string, selection: string, queryInput: QueryInput);
+	getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<ResultSetSubset | undefined>;
+	runQuery(uri: string, selection: ISelectionData | undefined, queryInput: QueryInput, runOptions?: ExecutionPlanOptions): void;
+	runQueryStatement(uri: string, selection: ISelectionData | undefined, queryInput: QueryInput): void;
+	runQueryString(uri: string, selection: string | undefined, queryInput: QueryInput): void;
 	cancelQuery(input: QueryRunner | string): void;
 	disposeQuery(uri: string): void;
 	isRunningQuery(uri: string): boolean;
@@ -80,16 +78,16 @@ export interface IQueryModelService {
 
 	// Edit Data Functions
 	initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): void;
-	disposeEdit(ownerUri: string): Thenable<void>;
-	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<EditUpdateCellResult>;
-	commitEdit(ownerUri): Thenable<void>;
-	createRow(ownerUri: string): Thenable<EditCreateRowResult>;
-	deleteRow(ownerUri: string, rowId: number): Thenable<void>;
-	revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<EditRevertCellResult>;
-	revertRow(ownerUri: string, rowId: number): Thenable<void>;
-	getEditRows(ownerUri: string, rowStart: number, numberOfRows: number): Thenable<EditSubsetResult>;
+	disposeEdit(ownerUri: string): Promise<void>;
+	updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<EditUpdateCellResult | undefined>;
+	commitEdit(ownerUri: string): Promise<void>;
+	createRow(ownerUri: string): Promise<EditCreateRowResult | undefined>;
+	deleteRow(ownerUri: string, rowId: number): Promise<void>;
+	revertCell(ownerUri: string, rowId: number, columnId: number): Promise<EditRevertCellResult | undefined>;
+	revertRow(ownerUri: string, rowId: number): Promise<void>;
+	getEditRows(ownerUri: string, rowStart: number, numberOfRows: number): Promise<EditSubsetResult | undefined>;
 
-	_getQueryInfo(uri: string): QueryInfo;
+	_getQueryInfo(uri: string): QueryInfo | undefined;
 	// Edit Data Callbacks
 	onEditSessionReady: Event<EditSessionReadyParams>;
 }

src/sql/platform/query/common/queryModelService.ts

@@ -36,7 +36,7 @@ export class QueryInfo {
 	public queryEventQueue: QueryEvent[];
 	public selection: Array<azdata.ISelectionData>;
 	public queryInput: QueryInput;
-	public selectionSnippet: string;
+	public selectionSnippet?: string;
 
 	// Notes if the angular components have obtained the DataService. If not, all messages sent
 	// via the data service will be lost.
@@ -85,7 +85,10 @@ export class QueryModelService implements IQueryModelService {
 
 	// IQUERYMODEL /////////////////////////////////////////////////////////
 	public getDataService(uri: string): DataService {
-		let dataService = this._getQueryInfo(uri).dataService;
+		let dataService: DataService | undefined;
+		if (this._queryInfoMap.has(uri)) {
+			dataService = this._getQueryInfo(uri)!.dataService;
+		}
 		if (!dataService) {
 			throw new Error('Could not find data service for uri: ' + uri);
 		}
@@ -119,40 +122,44 @@ export class QueryModelService implements IQueryModelService {
 	 * angular is listening for them.
 	 */
 	public onAngularLoaded(uri: string) {
-		let info = this._getQueryInfo(uri);
-		info.dataServiceReady = true;
-		this._sendQueuedEvents(uri);
+		if (this._queryInfoMap.has(uri)) {
+			let info = this._getQueryInfo(uri)!;
+			info.dataServiceReady = true;
+			this._sendQueuedEvents(uri);
+		}
 	}
 
 	/**
 	 * Get more data rows from the current resultSets from the service layer
 	 */
-	public getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Thenable<azdata.ResultSetSubset> {
-		return this._getQueryInfo(uri).queryRunner.getQueryRows(rowStart, numberOfRows, batchId, resultId).then(results => {
-			return results.resultSubset;
-		});
+	public getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<azdata.ResultSetSubset | undefined> {
+		if (this._queryInfoMap.has(uri)) {
+			return this._getQueryInfo(uri)!.queryRunner.getQueryRows(rowStart, numberOfRows, batchId, resultId).then(results => {
+				return results.resultSubset;
+			});
+		} else {
+			return Promise.resolve(undefined);
+		}
 	}
 
-	public getEditRows(uri: string, rowStart: number, numberOfRows: number): Thenable<azdata.EditSubsetResult> {
-		return this._queryInfoMap.get(uri).queryRunner.getEditRows(rowStart, numberOfRows).then(results => {
-			return results;
-		});
-	}
-
-	public getConfig(): Promise<{ [key: string]: any }> {
-		return undefined;
-	}
-
-	public getShortcuts(): Promise<any> {
-		return undefined;
+	public getEditRows(uri: string, rowStart: number, numberOfRows: number): Promise<azdata.EditSubsetResult | undefined> {
+		if (this._queryInfoMap.has(uri)) {
+			return this._queryInfoMap.get(uri)!.queryRunner.getEditRows(rowStart, numberOfRows).then(results => {
+				return results;
+			});
+		} else {
+			return Promise.resolve(undefined);
+		}
 	}
 
 	public copyResults(uri: string, selection: Slick.Range[], batchId: number, resultId: number, includeHeaders?: boolean): void {
-		this._queryInfoMap.get(uri).queryRunner.copyResults(selection, batchId, resultId, includeHeaders);
+		if (this._queryInfoMap.has(uri)) {
+			this._queryInfoMap.get(uri)!.queryRunner.copyResults(selection, batchId, resultId, includeHeaders);
+		}
 	}
 
 	public setEditorSelection(uri: string, index: number): void {
-		let info: QueryInfo = this._queryInfoMap.get(uri);
+		let info = this._queryInfoMap.get(uri);
 		if (info && info.queryInput) {
 			info.queryInput.updateSelection(info.selection[index]);
 		}
@@ -174,7 +181,7 @@ export class QueryModelService implements IQueryModelService {
 	public isRunningQuery(uri: string): boolean {
 		return !this._queryInfoMap.has(uri)
 			? false
-			: this._getQueryInfo(uri).queryRunner.isExecuting;
+			: this._getQueryInfo(uri)!.queryRunner.isExecuting;
 	}
 
 	/**
@@ -204,11 +211,11 @@ export class QueryModelService implements IQueryModelService {
 	private doRunQuery(uri: string, selection: azdata.ISelectionData | string, queryInput: QueryInput,
 		runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): void {
 		// Reuse existing query runner if it exists
-		let queryRunner: QueryRunner;
+		let queryRunner: QueryRunner | undefined;
 		let info: QueryInfo;
 
 		if (this._queryInfoMap.has(uri)) {
-			info = this._getQueryInfo(uri);
+			info = this._getQueryInfo(uri)!;
 			let existingRunner: QueryRunner = info.queryRunner;
 
 			// If the query is already in progress, don't attempt to send it
@@ -365,11 +372,11 @@ export class QueryModelService implements IQueryModelService {
 	}
 
 	public cancelQuery(input: QueryRunner | string): void {
-		let queryRunner: QueryRunner;
+		let queryRunner: QueryRunner | undefined;
 
 		if (typeof input === 'string') {
 			if (this._queryInfoMap.has(input)) {
-				queryRunner = this._getQueryInfo(input).queryRunner;
+				queryRunner = this._getQueryInfo(input)!.queryRunner;
 			}
 		} else {
 			queryRunner = input;
@@ -391,7 +398,7 @@ export class QueryModelService implements IQueryModelService {
 				severity: Severity.Error,
 				message: strings.format(LocalizedConstants.msgCancelQueryFailed, error)
 			});
-			this._fireQueryEvent(queryRunner.uri, 'complete', 0);
+			this._fireQueryEvent(queryRunner!.uri, 'complete', 0);
 		});
 
 	}
@@ -415,7 +422,7 @@ export class QueryModelService implements IQueryModelService {
 		let info: QueryInfo;
 
 		if (this._queryInfoMap.has(ownerUri)) {
-			info = this._getQueryInfo(ownerUri);
+			info = this._getQueryInfo(ownerUri)!;
 			let existingRunner: QueryRunner = info.queryRunner;
 
 			// If the initialization is already in progress
@@ -522,16 +529,16 @@ export class QueryModelService implements IQueryModelService {
 		// TODO: Implement query cancellation service
 	}
 
-	public disposeEdit(ownerUri: string): Thenable<void> {
+	public disposeEdit(ownerUri: string): Promise<void> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
 			return queryRunner.disposeEdit(ownerUri);
 		}
-		return Promise.resolve(null);
+		return Promise.resolve();
 	}
 
-	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<azdata.EditUpdateCellResult> {
+	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult | undefined> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
@@ -543,10 +550,10 @@ export class QueryModelService implements IQueryModelService {
 				return Promise.reject(error);
 			});
 		}
-		return Promise.resolve(null);
+		return Promise.resolve(undefined);
 	}
 
-	public commitEdit(ownerUri): Thenable<void> {
+	public commitEdit(ownerUri: string): Promise<void> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
@@ -558,49 +565,49 @@ export class QueryModelService implements IQueryModelService {
 				return Promise.reject(error);
 			});
 		}
-		return Promise.resolve(null);
+		return Promise.resolve();
 	}
 
-	public createRow(ownerUri: string): Thenable<azdata.EditCreateRowResult> {
+	public createRow(ownerUri: string): Promise<azdata.EditCreateRowResult | undefined> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
 			return queryRunner.createRow(ownerUri);
 		}
-		return Promise.resolve(null);
+		return Promise.resolve(undefined);
 	}
 
-	public deleteRow(ownerUri: string, rowId: number): Thenable<void> {
+	public deleteRow(ownerUri: string, rowId: number): Promise<void> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
 			return queryRunner.deleteRow(ownerUri, rowId);
 		}
-		return Promise.resolve(null);
+		return Promise.resolve();
 	}
 
-	public revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<azdata.EditRevertCellResult> {
+	public revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult | undefined> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
 			return queryRunner.revertCell(ownerUri, rowId, columnId);
 		}
-		return Promise.resolve(null);
+		return Promise.resolve(undefined);
 	}
 
-	public revertRow(ownerUri: string, rowId: number): Thenable<void> {
+	public revertRow(ownerUri: string, rowId: number): Promise<void> {
 		// Get existing query runner
 		let queryRunner = this.internalGetQueryRunner(ownerUri);
 		if (queryRunner) {
 			return queryRunner.revertRow(ownerUri, rowId);
 		}
-		return Promise.resolve(null);
+		return Promise.resolve();
 	}
 
-	public getQueryRunner(ownerUri): QueryRunner {
-		let queryRunner: QueryRunner = undefined;
+	public getQueryRunner(ownerUri: string): QueryRunner | undefined {
+		let queryRunner: QueryRunner | undefined = undefined;
 		if (this._queryInfoMap.has(ownerUri)) {
-			queryRunner = this._getQueryInfo(ownerUri).queryRunner;
+			queryRunner = this._getQueryInfo(ownerUri)!.queryRunner;
 		}
 		// return undefined if not found or is already executing
 		return queryRunner;
@@ -608,13 +615,13 @@ export class QueryModelService implements IQueryModelService {
 
 	// PRIVATE METHODS //////////////////////////////////////////////////////
 
-	private internalGetQueryRunner(ownerUri): QueryRunner {
-		let queryRunner: QueryRunner = undefined;
+	private internalGetQueryRunner(ownerUri: string): QueryRunner | undefined {
+		let queryRunner: QueryRunner | undefined;
 		if (this._queryInfoMap.has(ownerUri)) {
-			let existingRunner = this._getQueryInfo(ownerUri).queryRunner;
+			let existingRunner = this._getQueryInfo(ownerUri)!.queryRunner;
 			// If the query is not already executing then set it up
 			if (!existingRunner.isExecuting) {
-				queryRunner = this._getQueryInfo(ownerUri).queryRunner;
+				queryRunner = this._getQueryInfo(ownerUri)!.queryRunner;
 			}
 		}
 		// return undefined if not found or is already executing
@@ -622,7 +629,7 @@ export class QueryModelService implements IQueryModelService {
 	}
 
 	private _fireGridContentEvent(uri: string, type: string): void {
-		let info: QueryInfo = this._getQueryInfo(uri);
+		let info = this._getQueryInfo(uri);
 
 		if (info && info.dataServiceReady) {
 			let service: DataService = this.getDataService(uri);
@@ -635,29 +642,29 @@ export class QueryModelService implements IQueryModelService {
 	}
 
 	private _fireQueryEvent(uri: string, type: string, data?: any) {
-		let info: QueryInfo = this._getQueryInfo(uri);
+		let info = this._getQueryInfo(uri);
 
-		if (info.dataServiceReady) {
+		if (info && info.dataServiceReady) {
 			let service: DataService = this.getDataService(uri);
 			service.queryEventObserver.next({
 				type: type,
 				data: data
 			});
-		} else {
+		} else if (info) {
 			let queueItem: QueryEvent = { type: type, data: data };
 			info.queryEventQueue.push(queueItem);
 		}
 	}
 
 	private _sendQueuedEvents(uri: string): void {
-		let info: QueryInfo = this._getQueryInfo(uri);
-		while (info.queryEventQueue.length > 0) {
-			let event: QueryEvent = info.queryEventQueue.shift();
+		let info = this._getQueryInfo(uri);
+		while (info && info.queryEventQueue.length > 0) {
+			let event = info.queryEventQueue.shift()!;
 			this._fireQueryEvent(uri, event.type, event.data);
 		}
 	}
 
-	public _getQueryInfo(uri: string): QueryInfo {
+	public _getQueryInfo(uri: string): QueryInfo | undefined {
 		return this._queryInfoMap.get(uri);
 	}
 

src/sql/platform/query/common/queryRunner.ts

@@ -55,7 +55,7 @@ export default class QueryRunner extends Disposable {
 	private _isQueryPlan: boolean = false;
 	public get isQueryPlan(): boolean { return this._isQueryPlan; }
 	private _planXml = new Deferred<string>();
-	public get planXml(): Thenable<string> { return this._planXml.promise; }
+	public get planXml(): Promise<string> { return this._planXml.promise; }
 
 	private _onMessage = this._register(new Emitter<IQueryMessage>());
 	public get onMessage(): Event<IQueryMessage> { return this._onMessage.event; } // this is the only way typemoq can moq this... needs investigation @todo anthonydresser 5/2/2019
@@ -87,12 +87,12 @@ export default class QueryRunner extends Disposable {
 	private _onVisualize = this._register(new Emitter<azdata.ResultSetSummary>());
 	public readonly onVisualize = this._onVisualize.event;
 
-	private _queryStartTime: Date;
-	public get queryStartTime(): Date {
+	private _queryStartTime?: Date;
+	public get queryStartTime(): Date | undefined {
 		return this._queryStartTime;
 	}
-	private _queryEndTime: Date;
-	public get queryEndTime(): Date {
+	private _queryEndTime?: Date;
+	public get queryEndTime(): Date | undefined {
 		return this._queryEndTime;
 	}
 
@@ -135,7 +135,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Cancels the running query, if there is one
 	 */
-	public cancelQuery(): Thenable<azdata.QueryCancelResult> {
+	public cancelQuery(): Promise<azdata.QueryCancelResult> {
 		return this._queryManagementService.cancelQuery(this.uri);
 	}
 
@@ -143,21 +143,25 @@ export default class QueryRunner extends Disposable {
 	 * Runs the query with the provided query
 	 * @param input Query string to execute
 	 */
-	public runQuery(input: string, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
+	public runQuery(input: string, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
 	/**
 	 * Runs the query by pulling the query from the document using the provided selection data
 	 * @param input selection data
 	 */
-	public runQuery(input: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
-	public runQuery(input, runOptions?: azdata.ExecutionPlanOptions): Thenable<void> {
-		return this.doRunQuery(input, false, runOptions);
+	public runQuery(input: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	public runQuery(input: string | azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
+		if (types.isString(input)) {
+			return this.doRunQuery(input, false, runOptions);
+		} else {
+			return this.doRunQuery(input, false, runOptions);
+		}
 	}
 
 	/**
 	 * Runs the current SQL statement by pulling the query from the document using the provided selection data
 	 * @param input selection data
 	 */
-	public runQueryStatement(input: azdata.ISelectionData): Thenable<void> {
+	public runQueryStatement(input: azdata.ISelectionData): Promise<void> {
 		return this.doRunQuery(input, true);
 	}
 
@@ -165,11 +169,11 @@ export default class QueryRunner extends Disposable {
 	 * Implementation that runs the query with the provided query
 	 * @param input Query string to execute
 	 */
-	private doRunQuery(input: string, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
-	private doRunQuery(input: azdata.ISelectionData, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Thenable<void>;
-	private doRunQuery(input, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Thenable<void> {
+	private doRunQuery(input: string, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	private doRunQuery(input: azdata.ISelectionData, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void>;
+	private doRunQuery(input: string | azdata.ISelectionData, runCurrentStatement: boolean, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
 		if (this.isExecuting) {
-			return Promise.resolve(undefined);
+			return Promise.resolve();
 		}
 		this._planXml = new Deferred<string>();
 		this._batchSets = [];
@@ -177,7 +181,7 @@ export default class QueryRunner extends Disposable {
 		this._queryStartTime = undefined;
 		this._queryEndTime = undefined;
 		this._messages = [];
-		if (types.isObject(input) || types.isUndefinedOrNull(input)) {
+		if (isSelectionOrUndefined(input)) {
 			// Update internal state to show that we're executing the query
 			this._resultLineOffset = input ? input.startLine : 0;
 			this._resultColumnOffset = input ? input.startColumn : 0;
@@ -191,7 +195,7 @@ export default class QueryRunner extends Disposable {
 			return runCurrentStatement
 				? this._queryManagementService.runQueryStatement(this.uri, input.startLine, input.startColumn).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e))
 				: this._queryManagementService.runQuery(this.uri, input, runOptions).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
-		} else if (types.isString(input)) {
+		} else {
 			// Update internal state to show that we're executing the query
 			this._isExecuting = true;
 			this._totalElapsedMilliseconds = 0;
@@ -199,8 +203,6 @@ export default class QueryRunner extends Disposable {
 			this._onQueryStart.fire();
 
 			return this._queryManagementService.runQueryString(this.uri, input).then(() => this.handleSuccessRunQueryResult(), e => this.handleFailureRunQueryResult(e));
-		} else {
-			return Promise.reject('Unknown input');
 		}
 	}
 
@@ -329,7 +331,7 @@ export default class QueryRunner extends Disposable {
 				if (this._batchSets.length > 0) {
 					batchSet = this._batchSets[0];
 				} else {
-					batchSet = <azdata.BatchSummary>{
+					batchSet = <azdata.BatchSummary><unknown>{
 						id: 0,
 						selection: undefined,
 						hasError: false,
@@ -399,7 +401,7 @@ export default class QueryRunner extends Disposable {
 	 */
 	public handleMessage(obj: azdata.QueryExecuteMessageParams): void {
 		let message = obj.message;
-		message.time = new Date(message.time).toLocaleTimeString();
+		message.time = new Date(message.time!).toLocaleTimeString();
 		this._messages.push(message);
 
 		// Send the message to the results pane
@@ -409,7 +411,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Get more data rows from the current resultSets from the service layer
 	 */
-	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Thenable<azdata.QueryExecuteSubsetResult> {
+	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Promise<azdata.QueryExecuteSubsetResult> {
 		let rowData: azdata.QueryExecuteSubsetParams = <azdata.QueryExecuteSubsetParams>{
 			ownerUri: this.uri,
 			resultSetIndex: resultSetIndex,
@@ -430,7 +432,7 @@ export default class QueryRunner extends Disposable {
 	/*
 	 * Handle a session ready event for Edit Data
 	 */
-	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Thenable<void> {
+	public initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
 		// Update internal state to show that we're executing the query
 		this._isExecuting = true;
 		this._totalElapsedMilliseconds = 0;
@@ -454,7 +456,7 @@ export default class QueryRunner extends Disposable {
 	 * @param rowStart     The index of the row to start returning (inclusive)
 	 * @param numberOfRows The number of rows to return
 	 */
-	public getEditRows(rowStart: number, numberOfRows: number): Thenable<azdata.EditSubsetResult> {
+	public getEditRows(rowStart: number, numberOfRows: number): Promise<azdata.EditSubsetResult> {
 		const self = this;
 		let rowData: azdata.EditSubsetParams = {
 			ownerUri: this.uri,
@@ -488,35 +490,35 @@ export default class QueryRunner extends Disposable {
 		this._onEditSessionReady.fire({ ownerUri, success, message });
 	}
 
-	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Thenable<azdata.EditUpdateCellResult> {
+	public updateCell(ownerUri: string, rowId: number, columnId: number, newValue: string): Promise<azdata.EditUpdateCellResult> {
 		return this._queryManagementService.updateCell(ownerUri, rowId, columnId, newValue);
 	}
 
-	public commitEdit(ownerUri): Thenable<void> {
+	public commitEdit(ownerUri: string): Promise<void> {
 		return this._queryManagementService.commitEdit(ownerUri);
 	}
 
-	public createRow(ownerUri: string): Thenable<azdata.EditCreateRowResult> {
+	public createRow(ownerUri: string): Promise<azdata.EditCreateRowResult> {
 		return this._queryManagementService.createRow(ownerUri).then(result => {
 			return result;
 		});
 	}
 
-	public deleteRow(ownerUri: string, rowId: number): Thenable<void> {
+	public deleteRow(ownerUri: string, rowId: number): Promise<void> {
 		return this._queryManagementService.deleteRow(ownerUri, rowId);
 	}
 
-	public revertCell(ownerUri: string, rowId: number, columnId: number): Thenable<azdata.EditRevertCellResult> {
+	public revertCell(ownerUri: string, rowId: number, columnId: number): Promise<azdata.EditRevertCellResult> {
 		return this._queryManagementService.revertCell(ownerUri, rowId, columnId).then(result => {
 			return result;
 		});
 	}
 
-	public revertRow(ownerUri: string, rowId: number): Thenable<void> {
+	public revertRow(ownerUri: string, rowId: number): Promise<void> {
 		return this._queryManagementService.revertRow(ownerUri, rowId);
 	}
 
-	public disposeEdit(ownerUri: string): Thenable<void> {
+	public disposeEdit(ownerUri: string): Promise<void> {
 		return this._queryManagementService.disposeEdit(ownerUri);
 	}
 
@@ -530,7 +532,7 @@ export default class QueryRunner extends Disposable {
 	}
 
 	public dispose() {
-		this._batchSets = undefined;
+		this._batchSets = undefined!;
 		super.dispose();
 	}
 
@@ -551,8 +553,8 @@ export default class QueryRunner extends Disposable {
 	}
 
 
-	public getColumnHeaders(batchId: number, resultId: number, range: Slick.Range): string[] {
-		let headers: string[] = undefined;
+	public getColumnHeaders(batchId: number, resultId: number, range: Slick.Range): string[] | undefined {
+		let headers: string[] | undefined = undefined;
 		let batchSummary: azdata.BatchSummary = this._batchSets[batchId];
 		if (batchSummary !== undefined) {
 			let resultSetSummary = batchSummary.resultSetSummaries[resultId];
@@ -612,7 +614,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	) {
 	}
 
-	getRowData(rowStart: number, numberOfRows: number): Thenable<azdata.QueryExecuteSubsetResult> {
+	getRowData(rowStart: number, numberOfRows: number): Promise<azdata.QueryExecuteSubsetResult> {
 		return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId);
 	}
 
@@ -637,7 +639,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	shouldRemoveNewLines(): boolean {
 		return shouldRemoveNewLines(this._configurationService);
 	}
-	getColumnHeaders(range: Slick.Range): string[] {
+	getColumnHeaders(range: Slick.Range): string[] | undefined {
 		return this.queryRunner.getColumnHeaders(this.batchId, this.resultSetId, range);
 	}
 
@@ -645,7 +647,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 		return true;
 	}
 
-	serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable<void> {
+	serializeResults(format: SaveFormat, selection: Slick.Range[]): Promise<void> {
 		return this.queryRunner.serializeResults(this.batchId, this.resultSetId, format, selection);
 	}
 }
@@ -670,3 +672,7 @@ export function shouldRemoveNewLines(configurationService: IConfigurationService
 	let removeNewLines = configurationService.getValue<boolean>('sql.copyRemoveNewLine');
 	return !!removeNewLines;
 }
+
+function isSelectionOrUndefined(input: string | azdata.ISelectionData | undefined): input is azdata.ISelectionData | undefined {
+	return types.isObject(input) || types.isUndefinedOrNull(input);
+}