diff --git a/extensions/mssql/package.json b/extensions/mssql/package.json index d89762fb36..a8f74b24cd 100644 --- a/extensions/mssql/package.json +++ b/extensions/mssql/package.json @@ -577,6 +577,7 @@ "connectionProvider": { "providerId": "MSSQL", "displayName": "%mssql.provider.displayName%", + "isExecutionPlanProvider": true, "supportedExecutionPlanFileExtensions": [ "sqlplan" ], diff --git a/src/sql/platform/capabilities/common/capabilitiesService.ts b/src/sql/platform/capabilities/common/capabilitiesService.ts index 713ba92836..c457bbfbc0 100644 --- a/src/sql/platform/capabilities/common/capabilitiesService.ts +++ b/src/sql/platform/capabilities/common/capabilitiesService.ts @@ -56,6 +56,7 @@ export interface ConnectionProviderProperties { azureResource?: string; connectionOptions: azdata.ConnectionOption[]; isQueryProvider?: boolean; + isExecutionPlanProvider?: boolean; supportedExecutionPlanFileExtensions?: string[]; connectionStringOptions?: ConnectionStringOptions; } diff --git a/src/sql/workbench/services/executionPlan/common/executionPlanService.ts b/src/sql/workbench/services/executionPlan/common/executionPlanService.ts index bff601a4a5..9be8e27760 100644 --- a/src/sql/workbench/services/executionPlan/common/executionPlanService.ts +++ b/src/sql/workbench/services/executionPlan/common/executionPlanService.ts @@ -27,21 +27,59 @@ export class ExecutionPlanService implements IExecutionPlanService { this._providerRegisterEvent = this._onProviderRegister.event; } + private async ensureCapabilitiesRegistered(): Promise { + let providers = Object.keys(this._capabilitiesService.providers); + if (!providers) { + await new Promise(resolve => { + this._capabilitiesService.onCapabilitiesRegistered(e => { + resolve(); + }); + }); + } + } + + private async getExecutionPlanProvider(providerId: string): Promise { + await this.ensureCapabilitiesRegistered(); + const provider = this._capabilitiesService.providers[providerId]; + // Return undefined if the provider is not registered or it is not a execution plan provider. + if (!provider || !provider.connection?.isExecutionPlanProvider) { + return undefined; + } + + let handler = this._providers[providerId]; + if (handler) { + return handler; + } + // Ensure the extension is loaded + await this._extensionService.whenInstalledExtensionsRegistered(); + handler = this._providers[providerId]; + if (!handler) { + handler = await new Promise((resolve, reject) => { + this._providerRegisterEvent(e => { + if (e.id === providerId) { + resolve(e.provider); + } + }); + setTimeout(() => { + /** + * Handling a possible edge case where provider registered event + * might have been called before we await for it. + */ + resolve(this._providers[providerId]); + }, 30000); + }); + } + return handler; + } + /** * Runs the actions using the provider that supports the fileFormat provided. * @param fileFormat fileformat of the underlying execution plan file. It is used to get the provider that support it. * @param action executionPlanService action to be performed. */ private async _runAction(fileFormat: string, action: (handler: azdata.executionPlan.ExecutionPlanProvider) => Thenable): Promise { + await this.ensureCapabilitiesRegistered(); let providers = Object.keys(this._capabilitiesService.providers); - if (!providers) { - providers = await new Promise(resolve => { - this._capabilitiesService.onCapabilitiesRegistered(e => { - resolve(Object.keys(this._capabilitiesService.providers)); - }); - }); - } - let epProviders: string[] = []; for (let i = 0; i < providers.length; i++) { const providerCapabilities = this._capabilitiesService.getCapabilities(providers[i]); @@ -81,70 +119,7 @@ export class ExecutionPlanService implements IExecutionPlanService { if (!selectedProvider) { return Promise.reject(new Error(localize('providerIdNotValidError', "Valid provider is required in order to interact with ExecutionPlanService"))); } - await this._extensionService.whenInstalledExtensionsRegistered(); - let handler = this._providers[selectedProvider]; - if (!handler) { - handler = await new Promise((resolve, reject) => { - this._providerRegisterEvent(e => { - if (e.id === selectedProvider) { - resolve(e.provider); - } - }); - setTimeout(() => { - /** - * Handling a possible edge case where provider registered event - * might have been called before we await for it. - */ - resolve(this._providers[selectedProvider]); - }, 30000); - }); - } - if (handler) { - return Promise.resolve(action(handler)); - } else { - return Promise.reject(new Error(localize('noHandlerRegistered', "No valid execution plan handler is registered"))); - } - } - - /** - * Runs the action using the specified provider. - * @param providerId The provider ID that will be used to run an action. - * @param action executionPlanService action to be performed. - */ - private async _runActionForProvider(providerId: string, action: (handler: azdata.executionPlan.ExecutionPlanProvider) => Thenable): Promise { - let providers = Object.keys(this._capabilitiesService.providers); - if (!providers) { - providers = await new Promise(resolve => { - this._capabilitiesService.onCapabilitiesRegistered(e => { - resolve(Object.keys(this._capabilitiesService.providers)); - }); - }); - } - - const selectedProvider: string | undefined = providers.find(p => p === providerId); - if (!selectedProvider) { - return Promise.reject(new Error(localize('providerIdNotValidError', "Valid provider is required in order to interact with ExecutionPlanService"))); - } - - await this._extensionService.whenInstalledExtensionsRegistered(); - let handler = this._providers[selectedProvider]; - if (!handler) { - handler = await new Promise((resolve, reject) => { - this._providerRegisterEvent(e => { - if (e.id === selectedProvider) { - resolve(e.provider); - } - }); - setTimeout(() => { - /** - * Handling a possible edge case where provider registered event - * might have been called before we await for it. - */ - resolve(this._providers[selectedProvider]); - }, 30000); - }); - } - + const handler = await this.getExecutionPlanProvider(selectedProvider); if (handler) { return Promise.resolve(action(handler)); } else { @@ -175,10 +150,17 @@ export class ExecutionPlanService implements IExecutionPlanService { }); } - isExecutionPlan(providerId: string, value: string): Promise { - return this._runActionForProvider(providerId, (runner) => { - return runner.isExecutionPlan(value); - }); + async isExecutionPlan(providerId: string, value: string): Promise { + const handler = await this.getExecutionPlanProvider(providerId); + if (handler) { + return await handler.isExecutionPlan(value); + } else { + // If the provider is not an execution plan provider, then return false. + return { + isExecutionPlan: false, + queryExecutionPlanFileExtension: '' + }; + } } getSupportedExecutionPlanExtensionsForProvider(providerId: string): string[] | undefined {