diff --git a/src/sql/platform/query/common/queryManagement.ts b/src/sql/platform/query/common/queryManagement.ts index 4434b5b646..31ee02f900 100644 --- a/src/sql/platform/query/common/queryManagement.ts +++ b/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; @@ -87,6 +89,7 @@ export class QueryManagementService implements IQueryManagementService { public _serviceBrand: any; private _requestHandlers = new Map(); + private _onHandlerAddedEmitter = new Emitter(); // public for testing only public _queryRunners = new Map(); @@ -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 { + 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 = { diff --git a/src/sql/workbench/services/notebook/common/notebookServiceImpl.ts b/src/sql/workbench/services/notebook/common/notebookServiceImpl.ts index 68424a0678..67fac842e0 100644 --- a/src/sql/workbench/services/notebook/common/notebookServiceImpl.ts +++ b/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(OVERRIDE_EDITOR_THEMING_SETTING); if (overrideEditorSetting !== this._overrideEditorThemeSetting) {