fix issue caused by execution plan provider (#21533)

extensions/mssql/package.json

@@ -577,6 +577,7 @@
     "connectionProvider": {
       "providerId": "MSSQL",
       "displayName": "%mssql.provider.displayName%",
+      "isExecutionPlanProvider": true,
       "supportedExecutionPlanFileExtensions": [
         "sqlplan"
       ],

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;
 }

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<void> {
+		let providers = Object.keys(this._capabilitiesService.providers);
+		if (!providers) {
+			await new Promise<void>(resolve => {
+				this._capabilitiesService.onCapabilitiesRegistered(e => {
+					resolve();
+				});
+			});
+		}
+	}
+
+	private async getExecutionPlanProvider(providerId: string): Promise<azdata.executionPlan.ExecutionPlanProvider> {
+		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<T>(fileFormat: string, action: (handler: azdata.executionPlan.ExecutionPlanProvider) => Thenable<T>): Promise<T> {
+		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();
+		const handler = await this.getExecutionPlanProvider(selectedProvider);
-		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<T>(providerId: string, action: (handler: azdata.executionPlan.ExecutionPlanProvider) => Thenable<T>): Promise<T> {
-		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);
-			});
-		}
-
 		if (handler) {
 			return Promise.resolve(action(handler));
 		} else {
@@ -175,10 +150,17 @@ export class ExecutionPlanService implements IExecutionPlanService {
 		});
 	}
 
-	isExecutionPlan(providerId: string, value: string): Promise<azdata.executionPlan.IsExecutionPlanResult> {
+	async isExecutionPlan(providerId: string, value: string): Promise<azdata.executionPlan.IsExecutionPlanResult> {
-		return this._runActionForProvider(providerId, (runner) => {
+		const handler = await this.getExecutionPlanProvider(providerId);
-			return runner.isExecutionPlan(value);
+		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 {