Change SQL kernel to check queryManagementService instead of hardcoding (#4098)

* Change SQL kernel to check queryManagementService instead of hardcoding

* addressing PR comments

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

@@ -12,6 +12,7 @@ import { TPromise } from 'vs/base/common/winjs.base';
 import * as TelemetryKeys from 'sql/common/telemetryKeys';
 import * as TelemetryUtils from 'sql/common/telemetryUtilities';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
+import { Event, Emitter } from 'vs/base/common/event';
 
 export const SERVICE_ID = 'queryManagementService';
 
@@ -22,6 +23,7 @@ export interface IQueryManagementService {
 
 	addQueryRequestHandler(queryType: string, runner: IQueryRequestHandler): IDisposable;
 	isProviderRegistered(providerId: string): boolean;
+	getRegisteredProviders(): string[];
 	registerRunner(runner: QueryRunner, uri: string): void;
 
 	cancelQuery(ownerUri: string): Thenable<sqlops.QueryCancelResult>;
@@ -87,6 +89,7 @@ export class QueryManagementService implements IQueryManagementService {
 	public _serviceBrand: any;
 
 	private _requestHandlers = new Map<string, IQueryRequestHandler>();
+	private _onHandlerAddedEmitter = new Emitter<string>();
 	// public for testing only
 	public _queryRunners = new Map<string, QueryRunner>();
 
@@ -137,6 +140,7 @@ export class QueryManagementService implements IQueryManagementService {
 
 	public addQueryRequestHandler(queryType: string, handler: IQueryRequestHandler): IDisposable {
 		this._requestHandlers.set(queryType, handler);
+		this._onHandlerAddedEmitter.fire(queryType);
 
 		return {
 			dispose: () => {
@@ -144,11 +148,19 @@ export class QueryManagementService implements IQueryManagementService {
 		};
 	}
 
+	public get onHandlerAdded(): Event<string> {
+		return this._onHandlerAddedEmitter.event;
+	}
+
 	public isProviderRegistered(providerId: string): boolean {
 		let handler = this._requestHandlers.get(providerId);
 		return !!handler;
 	}
 
+	public getRegisteredProviders(): string[] {
+		return Array.from(this._requestHandlers.keys());
+	}
+
 	private addTelemetry(eventName: string, ownerUri: string, runOptions?: sqlops.ExecutionPlanOptions): void {
 		let providerId: string = this._connectionService.getProviderIdFromUri(ownerUri);
 		let data: TelemetryUtils.IConnectionTelemetryData = {

src/sql/workbench/services/notebook/common/notebookServiceImpl.ts

@@ -37,7 +37,8 @@ import { IEditorGroupsService } from 'vs/workbench/services/group/common/editorG
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IThemeService } from 'vs/platform/theme/common/themeService';
 import { registerNotebookThemes } from 'sql/parts/notebook/notebookStyles';
-import { ILanguageMagic } from 'sql/parts/notebook/models/modelInterfaces';
+import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
+import { ILanguageMagic, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
 
 export interface NotebookProviderProperties {
 	provider: string;
@@ -104,7 +105,8 @@ export class NotebookService extends Disposable implements INotebookService {
 		@IEditorService private readonly _editorService: IEditorService,
 		@IEditorGroupsService private readonly _editorGroupsService: IEditorGroupsService,
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
-		@IThemeService private readonly _themeService: IThemeService
+		@IThemeService private readonly _themeService: IThemeService,
+		@IQueryManagementService private readonly _queryManagementService
 	) {
 		super();
 		this._register(notebookRegistry.onNewRegistration(this.updateRegisteredProviders, this));
@@ -113,8 +115,15 @@ export class NotebookService extends Disposable implements INotebookService {
 		if (extensionService) {
 			extensionService.whenInstalledExtensionsRegistered().then(() => {
 				this.cleanupProviders();
-				this._isRegistrationComplete = true;
+
-				this._registrationComplete.resolve();
+				// If providers have already registered by this point, add them now (since onHandlerAdded will never fire)
+				if (this._queryManagementService.registeredProviders && this._queryManagementService.registeredProviders.length > 0) {
+					this.updateSQLRegistrationWithConnectionProviders();
+				}
+
+				this._register(this._queryManagementService.onHandlerAdded((queryType) => {
+					this.updateSQLRegistrationWithConnectionProviders();
+				}));
 			});
 		}
 		if (extensionManagementService) {
@@ -158,6 +167,23 @@ export class NotebookService extends Disposable implements INotebookService {
 		}
 	}
 
+	private updateSQLRegistrationWithConnectionProviders() {
+		// Update the SQL extension
+		let sqlNotebookProvider = this._providerToStandardKernels.get(notebookConstants.SQL);
+		if (sqlNotebookProvider) {
+			let sqlConnectionTypes = this._queryManagementService.getRegisteredProviders();
+			let provider = sqlNotebookProvider.find(p => p.name === notebookConstants.SQL);
+			if (provider) {
+				this._providerToStandardKernels.set(notebookConstants.SQL, [{
+					name: notebookConstants.SQL,
+					connectionProviderIds: sqlConnectionTypes
+				}]);
+			}
+		}
+		this._isRegistrationComplete = true;
+		this._registrationComplete.resolve();
+	}
+
 	private updateNotebookThemes() {
 		let overrideEditorSetting = this._configurationService.getValue<boolean>(OVERRIDE_EDITOR_THEMING_SETTING);
 		if (overrideEditorSetting !== this._overrideEditorThemeSetting) {