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>

extensions/mssql/src/contracts.ts

@@ -1560,6 +1560,22 @@ export namespace ExecutionPlanComparisonRequest {
 
 // ------------------------------- < Execution Plan > ------------------------------------
 
+// ------------------------------- < Server Contextualization API > ------------------------------------
+
+export interface ServerContextualizationParams {
+	ownerUri: string;
+}
+
+export namespace GenerateServerContextualizationNotification {
+	export const type = new NotificationType<ServerContextualizationParams, void>('metadata/generateServerContext');
+}
+
+export namespace GetServerContextualizationRequest {
+	export const type = new RequestType<ServerContextualizationParams, azdata.contextualization.GetServerContextualizationResult, void, void>('metadata/getServerContext');
+}
+
+// ------------------------------- < Database Server Contextualization API > ------------------------------------
+
 // ------------------------------- < Object Management > ------------------------------------
 export interface InitializeViewRequestParams {
 	connectionUri: string;

extensions/mssql/src/features.ts

@@ -1304,3 +1304,58 @@ export class ExecutionPlanServiceFeature extends SqlOpsFeature<undefined> {
 		});
 	}
 }
+
+/**
+ * Server Contextualization Service Feature
+ */
+export class ServerContextualizationServiceFeature extends SqlOpsFeature<undefined> {
+	private static readonly messagesTypes: RPCMessageType[] = [
+		contracts.GenerateServerContextualizationNotification.type
+	];
+
+	constructor(client: SqlOpsDataClient) {
+		super(client, ServerContextualizationServiceFeature.messagesTypes);
+	}
+
+	public fillClientCapabilities(capabilities: ClientCapabilities): void {
+	}
+
+	public initialize(capabilities: ServerCapabilities): void {
+		this.register(this.messages, {
+			id: UUID.generateUuid(),
+			registerOptions: undefined
+		});
+	}
+
+	protected registerProvider(options: undefined): Disposable {
+		const client = this._client;
+
+		const generateServerContextualization = (ownerUri: string): void => {
+			const params: contracts.ServerContextualizationParams = {
+				ownerUri: ownerUri
+			};
+
+			return client.sendNotification(contracts.GenerateServerContextualizationNotification.type, params);
+		};
+
+		const getServerContextualization = (ownerUri: string): Thenable<azdata.contextualization.GetServerContextualizationResult> => {
+			const params: contracts.ServerContextualizationParams = {
+				ownerUri: ownerUri
+			};
+
+			return client.sendRequest(contracts.GetServerContextualizationRequest.type, params).then(
+				r => r,
+				e => {
+					client.logFailedRequest(contracts.GetServerContextualizationRequest.type, e);
+					return Promise.reject(e);
+				}
+			);
+		};
+
+		return azdata.dataprotocol.registerServerContextualizationProvider({
+			providerId: client.providerId,
+			generateServerContextualization: generateServerContextualization,
+			getServerContextualization: getServerContextualization
+		});
+	}
+}

extensions/mssql/src/sqlToolsServer.ts

@@ -13,7 +13,7 @@ import * as azurecore from 'azurecore';
 import { getCommonLaunchArgsAndCleanupOldLogFiles, getConfigTracingLevel, getEnableConnectionPoolingConfig, getEnableSqlAuthenticationProviderConfig, getOrDownloadServer, getParallelMessageProcessingConfig, logDebug, TracingLevel } from './utils';
 import { TelemetryReporter, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
-import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature, TableDesignerFeature, ExecutionPlanServiceFeature } from './features';
+import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature, TableDesignerFeature, ExecutionPlanServiceFeature/*, ServerContextualizationServiceFeature*/ } from './features'; // LEWISSANCHEZ TODO: Put back ServerContextualizationServiceFeature once ready.
 import { CredentialStore } from './credentialstore/credentialstore';
 import { AzureResourceProvider } from './resourceProvider/resourceProvider';
 import { SchemaCompareService } from './schemaCompare/schemaCompareService';
@@ -247,6 +247,7 @@ function getClientOptions(context: AppContext): ClientOptions {
 			SqlCredentialService.asFeature(context),
 			TableDesignerFeature,
 			ExecutionPlanServiceFeature,
+			// ServerContextualizationServiceFeature, // LEWISSANCHEZ TODO: Put this provider back once STS changes are complete
 			ErrorDiagnosticsProvider.asFeature(context),
 			ObjectManagementService.asFeature(context)
 		],