Adds functionality to open execution plans from DB tables and DMVs. (#20377)

* Adds functionality to open execution plans from DB tables and DMVs.

* Code review changes

* Renames method to getProviderFromUri

extensions/mssql/src/features.ts

@@ -1279,10 +1279,30 @@ export class ExecutionPlanServiceFeature extends SqlOpsFeature<undefined> {
 			);
 		};
 
+		const isExecutionPlan = (value: string): Thenable<azdata.executionPlan.IsExecutionPlanResult> => {
+			return new Promise((resolve) => {
+				let isExecutionPlan = false;
+				let queryExecutionPlanFileExtension = '';
+
+				if (value.includes('ShowPlanXML')) {
+					isExecutionPlan = true;
+					queryExecutionPlanFileExtension = 'sqlplan';
+				}
+
+				const result: azdata.executionPlan.IsExecutionPlanResult = {
+					isExecutionPlan: isExecutionPlan,
+					queryExecutionPlanFileExtension: queryExecutionPlanFileExtension,
+				};
+
+				return resolve(result);
+			});
+		};
+
 		return azdata.dataprotocol.registerExecutionPlanProvider({
 			providerId: client.providerId,
 			getExecutionPlan,
-			compareExecutionPlanGraph
+			compareExecutionPlanGraph,
+			isExecutionPlan
 		});
 	}
 }

src/sql/azdata.proposed.d.ts

@@ -1464,6 +1464,11 @@ declare module 'azdata' {
 			secondComparisonResult: ExecutionGraphComparisonResult;
 		}
 
+		export interface IsExecutionPlanResult {
+			isExecutionPlan: boolean;
+			queryExecutionPlanFileExtension: string;
+		}
+
 		export interface ExecutionPlanProvider extends DataProvider {
 			// execution plan service methods
 
@@ -1478,6 +1483,11 @@ declare module 'azdata' {
 			 * @param secondPlanFile file that contains the second execution plan.
 			 */
 			compareExecutionPlanGraph(firstPlanFile: ExecutionPlanGraphInfo, secondPlanFile: ExecutionPlanGraphInfo): Thenable<ExecutionPlanComparisonResult>;
+			/**
+			 * Determines if the provided value is an execution plan and returns the appropriate file extension.
+			 * @param value String that needs to be checked.
+			 */
+			isExecutionPlan(value: string): Thenable<IsExecutionPlanResult>;
 		}
 
 		export interface TopOperationsDataItem {

src/sql/workbench/api/browser/mainThreadDataProtocol.ts

@@ -554,7 +554,8 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData
 	public $registerExecutionPlanProvider(providerId: string, handle: number): void {
 		this._executionPlanService.registerProvider(providerId, <azdata.executionPlan.ExecutionPlanProvider>{
 			getExecutionPlan: (planFile: azdata.executionPlan.ExecutionPlanGraphInfo) => this._proxy.$getExecutionPlan(handle, planFile),
-			compareExecutionPlanGraph: (firstPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo, secondPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo) => this._proxy.$compareExecutionPlanGraph(handle, firstPlanFile, secondPlanFile)
+			compareExecutionPlanGraph: (firstPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo, secondPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo) => this._proxy.$compareExecutionPlanGraph(handle, firstPlanFile, secondPlanFile),
+			isExecutionPlan: (value: string) => this._proxy.$isExecutionPlan(handle, value)
 		});
 	}
 

src/sql/workbench/api/common/extHostDataProtocol.ts

@@ -937,4 +937,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	public override $compareExecutionPlanGraph(handle: number, firstPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo, secondPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo): Thenable<azdata.executionPlan.ExecutionPlanComparisonResult> {
 		return this._resolveProvider<azdata.executionPlan.ExecutionPlanProvider>(handle).compareExecutionPlanGraph(firstPlanFile, secondPlanFile);
 	}
+
+	public override $isExecutionPlan(handle: number, value: string): Thenable<azdata.executionPlan.IsExecutionPlanResult> {
+		return this._resolveProvider<azdata.executionPlan.ExecutionPlanProvider>(handle).isExecutionPlan(value);
+	}
 }

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -585,6 +585,10 @@ export abstract class ExtHostDataProtocolShape {
 	 * Compares two execution plans and identifies matching sections in both.
 	 */
 	$compareExecutionPlanGraph(handle: number, firstPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo, secondPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo): Thenable<azdata.executionPlan.ExecutionPlanComparisonResult> { throw ni(); }
+	/**
+	 * Determines if the provided value is an execution plan and returns the appropriate file extension.
+	 */
+	$isExecutionPlan(handle: number, value: string): Thenable<azdata.executionPlan.IsExecutionPlanResult> { throw ni(); }
 }
 
 /**

src/sql/workbench/contrib/executionPlan/browser/executionPlanContribution.ts

@@ -70,7 +70,12 @@ export class ExecutionPlanEditorOverrideContribution extends Disposable implemen
 			},
 			{},
 			(editorInput, group) => {
-				const executionPlanInput = this._instantiationService.createInstance(ExecutionPlanInput, editorInput.resource);
+				const executionPlanGraphInfo = {
+					graphFileContent: undefined,
+					graphFileType: undefined
+				};
+
+				const executionPlanInput = this._instantiationService.createInstance(ExecutionPlanInput, editorInput.resource, executionPlanGraphInfo);
 				return { editor: executionPlanInput, options: editorInput.options, group: group };
 			}
 		);

src/sql/workbench/contrib/executionPlan/common/executionPlanInput.ts

@@ -5,23 +5,41 @@
 
 import * as path from 'vs/base/common/path';
 import { URI } from 'vs/base/common/uri';
+import { localize } from 'vs/nls';
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { EditorModel } from 'vs/workbench/common/editor/editorModel';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { ITextFileService } from 'vs/workbench/services/textfile/common/textfiles';
+import * as azdata from 'azdata';
 
 export class ExecutionPlanInput extends EditorInput {
 
 	public static ID: string = 'workbench.editorinputs.executionplan';
 	public static SCHEMA: string = 'executionplan';
+	private readonly editorNamePrefix = localize('epCompare.executionPlanEditorName', 'ExecutionPlan');
+	private _editorName: string;
 
-	private _content?: string;
 	public _executionPlanFileViewUUID: string;
 
 	constructor(
-		private _uri: URI,
+		private _uri: URI | undefined,
+		private executionPlanGraphinfo: azdata.executionPlan.ExecutionPlanGraphInfo,
 		@ITextFileService private readonly _fileService: ITextFileService,
+		@IEditorService private readonly _editorService: IEditorService
 	) {
 		super();
+
+		if (this._uri === undefined && !!this.executionPlanGraphinfo.graphFileContent && !!this.executionPlanGraphinfo.graphFileType) {
+			const existingNames = this._editorService.editors.map(editor => editor.getName());
+			let i = 0;
+			this._editorName = `${this.editorNamePrefix}${i}.${this.executionPlanGraphinfo.graphFileType}`;
+			while (existingNames.includes(this._editorName)) {
+				i++;
+				this._editorName = `${this.editorNamePrefix}${i}.${this.executionPlanGraphinfo.graphFileType}`;
+			}
+
+			this._uri = URI.parse(this._editorName);
+		}
 	}
 
 	public get executionPlanFileViewUUID(): string {
@@ -37,14 +55,18 @@ export class ExecutionPlanInput extends EditorInput {
 	}
 
 	public override getName(): string {
+		if (this._editorName) {
+			return this._editorName;
+		}
+
 		return path.basename(this._uri.fsPath);
 	}
 
 	public async content(): Promise<string> {
-		if (!this._content) {
+		if (!this.executionPlanGraphinfo.graphFileContent) {
-			this._content = (await this._fileService.read(this._uri, { acceptTextOnly: true })).value;
+			this.executionPlanGraphinfo.graphFileContent = (await this._fileService.read(this._uri, { acceptTextOnly: true })).value;
 		}
-		return this._content;
+		return this.executionPlanGraphinfo.graphFileContent;
 	}
 
 	public getUri(): string {
@@ -60,8 +82,8 @@ export class ExecutionPlanInput extends EditorInput {
 	}
 
 	public override async resolve(refresh?: boolean): Promise<EditorModel | undefined> {
-		if (!this._content) {
+		if (!this.executionPlanGraphinfo.graphFileContent) {
-			this._content = (await this._fileService.read(this._uri, { acceptTextOnly: true })).value;
+			this.executionPlanGraphinfo.graphFileContent = (await this._fileService.read(this._uri, { acceptTextOnly: true })).value;
 		}
 		return undefined;
 	}

src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts

@@ -47,6 +47,7 @@ import { IStorageService } from 'vs/platform/storage/common/storage';
 import { getChartMaxRowCount, notifyMaxRowCountExceeded } from 'sql/workbench/contrib/charts/browser/utils';
 import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
+import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces';
 
 @Component({
 	selector: GridOutputComponent.SELECTOR,
@@ -237,12 +238,13 @@ class DataResourceTable extends GridTableBase<any> {
 		@IQueryModelService queryModelService: IQueryModelService,
 		@IThemeService themeService: IThemeService,
 		@IContextViewService contextViewService: IContextViewService,
-		@INotificationService notificationService: INotificationService
+		@INotificationService notificationService: INotificationService,
+		@IExecutionPlanService executionPlanService: IExecutionPlanService
 	) {
 		super(state, createResultSet(source), {
 			actionOrientation: ActionsOrientation.HORIZONTAL,
 			inMemoryDataProcessing: true
-		}, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService, queryModelService, themeService, contextViewService, notificationService);
+		}, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService, queryModelService, themeService, contextViewService, notificationService, executionPlanService);
 		this._gridDataProvider = this.instantiationService.createInstance(DataResourceDataProvider, source, this.resultSet, this.cellModel);
 		this._chart = this.instantiationService.createInstance(ChartView, false);
 

src/sql/workbench/contrib/query/browser/gridPanel.ts

@@ -52,6 +52,8 @@ import { HybridDataProvider } from 'sql/base/browser/ui/table/hybridDataProvider
 import { INotificationService } from 'vs/platform/notification/common/notification';
 import { alert, status } from 'vs/base/browser/ui/aria/aria';
 import { CopyAction } from 'vs/editor/contrib/clipboard/clipboard';
+import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces';
+import { ExecutionPlanInput } from 'sql/workbench/contrib/executionPlan/common/executionPlanInput';
 
 const ROW_HEIGHT = 29;
 const HEADER_HEIGHT = 26;
@@ -375,6 +377,8 @@ export abstract class GridTableBase<T> extends Disposable implements IView {
 
 	public isOnlyTable: boolean = true;
 
+	public providerId: string;
+
 	// this handles if the row count is small, like 4-5 rows
 	protected get maxSize(): number {
 		return ((this.resultSet.rowCount) * this.rowHeight) + HEADER_HEIGHT + ESTIMATED_SCROLL_BAR_HEIGHT;
@@ -400,7 +404,8 @@ export abstract class GridTableBase<T> extends Disposable implements IView {
 		@IQueryModelService private readonly queryModelService: IQueryModelService,
 		@IThemeService private readonly themeService: IThemeService,
 		@IContextViewService private readonly contextViewService: IContextViewService,
-		@INotificationService private readonly notificationService: INotificationService
+		@INotificationService private readonly notificationService: INotificationService,
+		@IExecutionPlanService private readonly executionPlanService: IExecutionPlanService
 	) {
 		super();
 
@@ -712,12 +717,24 @@ export abstract class GridTableBase<T> extends Disposable implements IView {
 			const value = subset[0][event.cell.cell - 1];
 			const isJson = isJsonCell(value);
 			if (column.isXml || isJson) {
-				const content = value.displayValue;
+				const result = await this.executionPlanService.isExecutionPlan(this.providerId, value.displayValue);
-				const input = this.untitledEditorService.create({ mode: column.isXml ? 'xml' : 'json', initialValue: content });
+				if (result.isExecutionPlan) {
-				await input.resolve();
+					const executionPlanGraphInfo = {
-				await this.instantiationService.invokeFunction(formatDocumentWithSelectedProvider, input.textEditorModel, FormattingMode.Explicit, Progress.None, CancellationToken.None);
+						graphFileContent: value.displayValue,
-				input.setDirty(false);
+						graphFileType: result.queryExecutionPlanFileExtension
-				await this.editorService.openEditor(input);
+					};
+
+					const executionPlanInput = this.instantiationService.createInstance(ExecutionPlanInput, undefined, executionPlanGraphInfo);
+					await this.editorService.openEditor(executionPlanInput);
+				}
+				else {
+					const content = value.displayValue;
+					const input = this.untitledEditorService.create({ mode: column.isXml ? 'xml' : 'json', initialValue: content });
+					await input.resolve();
+					await this.instantiationService.invokeFunction(formatDocumentWithSelectedProvider, input.textEditorModel, FormattingMode.Explicit, Progress.None, CancellationToken.None);
+					input.setDirty(false);
+					await this.editorService.openEditor(input);
+				}
 			}
 		}
 	}
@@ -901,15 +918,17 @@ class GridTable<T> extends GridTableBase<T> {
 		@IQueryModelService queryModelService: IQueryModelService,
 		@IThemeService themeService: IThemeService,
 		@IContextViewService contextViewService: IContextViewService,
-		@INotificationService notificationService: INotificationService
+		@INotificationService notificationService: INotificationService,
+		@IExecutionPlanService executionPlanService: IExecutionPlanService
 	) {
 		super(state, resultSet, {
 			actionOrientation: ActionsOrientation.VERTICAL,
 			inMemoryDataProcessing: true,
 			showActionBar: true,
 			inMemoryDataCountThreshold: configurationService.getValue<IQueryEditorConfiguration>('queryEditor').results.inMemoryDataProcessingThreshold,
-		}, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService, queryModelService, themeService, contextViewService, notificationService);
+		}, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService, queryModelService, themeService, contextViewService, notificationService, executionPlanService);
 		this._gridDataProvider = this.instantiationService.createInstance(QueryGridDataProvider, this._runner, resultSet.batchId, resultSet.id);
+		this.providerId = this._runner.getProviderId();
 	}
 
 	get gridDataProvider(): IGridDataProvider {

src/sql/workbench/services/executionPlan/common/executionPlanService.ts

@@ -106,6 +106,52 @@ export class ExecutionPlanService implements IExecutionPlanService {
 		}
 	}
 
+	/**
+	 * 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 {
+			return Promise.reject(new Error(localize('noHandlerRegistered', "No valid execution plan handler is registered")));
+		}
+	}
+
 	registerProvider(providerId: string, provider: azdata.executionPlan.ExecutionPlanProvider): void {
 		if (this._providers[providerId]) {
 			throw new Error(`A execution plan provider with id "${providerId}" is already registered`);
@@ -129,6 +175,12 @@ export class ExecutionPlanService implements IExecutionPlanService {
 		});
 	}
 
+	isExecutionPlan(providerId: string, value: string): Promise<azdata.executionPlan.IsExecutionPlanResult> {
+		return this._runActionForProvider(providerId, (runner) => {
+			return runner.isExecutionPlan(value);
+		});
+	}
+
 	getSupportedExecutionPlanExtensionsForProvider(providerId: string): string[] | undefined {
 		return this._capabilitiesService.getCapabilities(providerId).connection.supportedExecutionPlanFileExtensions;
 	}

src/sql/workbench/services/executionPlan/common/interfaces.ts

@@ -28,6 +28,11 @@ export interface IExecutionPlanService {
 	 * @param secondPlanFile file that contains the second execution plan.
 	 */
 	compareExecutionPlanGraph(firstPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo, secondPlanFile: azdata.executionPlan.ExecutionPlanGraphInfo): Promise<azdata.executionPlan.ExecutionPlanComparisonResult>;
+	/**
+	 * Determines if the provided value is an execution plan and returns the appropriate file extension.
+	 * @param value String that needs to be checked.
+	 */
+	isExecutionPlan(providerId: string, value: string): Promise<azdata.executionPlan.IsExecutionPlanResult>;
 
 	/**
 	 * Get execution plan file extensions supported by all registered providers.

src/sql/workbench/services/query/common/queryManagement.ts

@@ -42,6 +42,7 @@ export interface IQueryManagementService {
 	getRegisteredProviders(): string[];
 	registerRunner(runner: QueryRunner, uri: string): void;
 	getRunner(uri: string): QueryRunner | undefined;
+	getProviderIdFromUri(uri: string): string;
 
 	cancelQuery(ownerUri: string): Promise<QueryCancelResult>;
 	runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void>;
@@ -149,6 +150,10 @@ export class QueryManagementService implements IQueryManagementService {
 		return this._queryRunners.get(uri);
 	}
 
+	public getProviderIdFromUri(uri: string): string {
+		return this._connectionService.getProviderIdFromUri(uri);
+	}
+
 	// Handles logic to run the given handlerCallback at the appropriate time. If the given runner is
 	// undefined, the handlerCallback is put on the _handlerCallbackQueue to be run once the runner is set
 	// public for testing only

src/sql/workbench/services/query/common/queryRunner.ts

@@ -128,6 +128,10 @@ export default class QueryRunner extends Disposable {
 		return this._messages.slice(0);
 	}
 
+	public getProviderId(): string {
+		return this.queryManagementService.getProviderIdFromUri(this.uri);
+	}
+
 	// PUBLIC METHODS ======================================================
 
 	/**

src/sql/workbench/services/query/test/common/testQueryManagementService.ts

@@ -29,6 +29,9 @@ export class TestQueryManagementService implements IQueryManagementService {
 	getRunner(uri: string): QueryRunner {
 		throw new Error('Method not implemented.');
 	}
+	getProviderIdFromUri(uri: string): string {
+		throw new Error('Method not implemented.');
+	}
 	async cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult> {
 		return { messages: undefined };
 	}