Integrate Contextualization API into Azure Data Studio to get better query recommendations from Copilot (#24044)

* Boilerplate for new metadata API endpoint * Register all server metadata provider * Fully registers data provider * Registers metadata provider * Instantiate metadata service * Generates server metadata when connection is established * Allow queryEditorInput to get server metadata * Minor clean up * Renames metadata provider and request endpoint * Corrects documentation block * Integrates get server metadata request endpoint * Adjusts GetServerMetadataResult scripts type * Add back Cargo.toml file * Fix SQL hygiene error * reflect changes made in in STS for table metadata * Adds feature toggle to serverMetadataService * Places toggle before request to get create scripts * Fix build check issues * Minor review changes * Improves contextualization setting label * Generalize contextualization service names * Additional code review changes * Update extensions/mssql/src/contracts.ts Co-authored-by: Charles Gagnon <chgagnon@microsoft.com> * Update src/sql/azdata.proposed.d.ts Co-authored-by: Charles Gagnon <chgagnon@microsoft.com> * Code reivew changes * Capitalize c in contextualization * Additional review changes * Update provider type * Simplify type and method names * Unregister MSSQL ServerContextualization provider --------- Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>
2026-01-25 09:35:37 -05:00 · 2023-08-21 19:54:44 -07:00
parent df5693ffd3
commit d5a9c172d1
19 changed files with 336 additions and 10 deletions
--- a/src/sql/workbench/services/contextualization/common/serverContextualizationService.ts
+++ b/src/sql/workbench/services/contextualization/common/serverContextualizationService.ts
@@ -0,0 +1,94 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import { invalidProvider } from 'sql/base/common/errors';
+import { IConnectionManagementService, IConnectionParams } from 'sql/platform/connection/common/connectionManagement';
+import { IQueryEditorConfiguration } from 'sql/platform/query/common/query';
+import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces';
+import { Disposable } from 'vs/base/common/lifecycle';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
+
+export class ServerContextualizationService extends Disposable implements IServerContextualizationService {
+	public _serviceBrand: undefined;
+	private _providers = new Map<string, azdata.contextualization.ServerContextualizationProvider>();
+
+	constructor(
+		@IConnectionManagementService private readonly _connectionManagementService: IConnectionManagementService,
+		@IConfigurationService private readonly _configurationService: IConfigurationService,
+		@IExtensionService private readonly _extensionService: IExtensionService
+	) {
+		super();
+
+		this._register(this._connectionManagementService.onConnect(async (e: IConnectionParams) => {
+			const copilotExt = await this._extensionService.getExtension('github.copilot');
+
+			if (copilotExt && this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled) {
+				const ownerUri = e.connectionUri;
+				await this.generateServerContextualization(ownerUri);
+			}
+		}));
+	}
+
+	/**
+	 * Register a server contextualization service provider
+	 */
+	public registerProvider(providerId: string, provider: azdata.contextualization.ServerContextualizationProvider): void {
+		if (this._providers.has(providerId)) {
+			throw new Error(`A server contextualization provider with ID "${providerId}" is already registered`);
+		}
+		this._providers.set(providerId, provider);
+	}
+
+	/**
+	 * Unregister a server contextualization service provider.
+	 */
+	public unregisterProvider(providerId: string): void {
+		this._providers.delete(providerId);
+	}
+
+	/**
+	 * Gets a registered server contextualization service provider. An exception is thrown if a provider isn't registered with the specified ID.
+	 * @param providerId The ID of the registered provider.
+	 */
+	public getProvider(providerId: string): azdata.contextualization.ServerContextualizationProvider {
+		const provider = this._providers.get(providerId);
+		if (provider) {
+			return provider;
+		}
+
+		throw invalidProvider(providerId);
+	}
+
+	/**
+	 * Generates server context
+	 * @param ownerUri The URI of the connection to generate context for.
+	 */
+	public generateServerContextualization(ownerUri: string): void {
+		const providerName = this._connectionManagementService.getProviderIdFromUri(ownerUri);
+		const handler = this.getProvider(providerName);
+		if (handler) {
+			handler.generateServerContextualization(ownerUri);
+		}
+	}
+
+	/**
+	 * Gets all database context.
+	 * @param ownerUri The URI of the connection to get context for.
+	 */
+	public async getServerContextualization(ownerUri: string): Promise<azdata.contextualization.GetServerContextualizationResult> {
+		const providerName = this._connectionManagementService.getProviderIdFromUri(ownerUri);
+		const handler = this.getProvider(providerName);
+		if (handler) {
+			return await handler.getServerContextualization(ownerUri);
+		}
+		else {
+			return Promise.resolve({
+				context: []
+			});
+		}
+	}
+}