diff --git a/extensions/mssql/src/contracts.ts b/extensions/mssql/src/contracts.ts index 2c77727c50..99cc33d9c1 100644 --- a/extensions/mssql/src/contracts.ts +++ b/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('metadata/generateServerContext'); +} + +export namespace GetServerContextualizationRequest { + export const type = new RequestType('metadata/getServerContext'); +} + +// ------------------------------- < Database Server Contextualization API > ------------------------------------ + // ------------------------------- < Object Management > ------------------------------------ export interface InitializeViewRequestParams { connectionUri: string; diff --git a/extensions/mssql/src/features.ts b/extensions/mssql/src/features.ts index ac318e67d3..9d79b97e1b 100644 --- a/extensions/mssql/src/features.ts +++ b/extensions/mssql/src/features.ts @@ -1304,3 +1304,58 @@ export class ExecutionPlanServiceFeature extends SqlOpsFeature { }); } } + +/** + * Server Contextualization Service Feature + */ +export class ServerContextualizationServiceFeature extends SqlOpsFeature { + 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 => { + 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 + }); + } +} diff --git a/extensions/mssql/src/sqlToolsServer.ts b/extensions/mssql/src/sqlToolsServer.ts index b4513786c2..f7f5cb9134 100644 --- a/extensions/mssql/src/sqlToolsServer.ts +++ b/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) ], diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index dccc39aeb2..07d3bbcf39 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/src/sql/azdata.proposed.d.ts @@ -887,12 +887,19 @@ declare module 'azdata' { export enum DataProviderType { TableDesignerProvider = 'TableDesignerProvider', - ExecutionPlanProvider = 'ExecutionPlanProvider' + ExecutionPlanProvider = 'ExecutionPlanProvider', + ServerContextualizationProvider = 'ServerContextualizationProvider' } export namespace dataprotocol { export function registerTableDesignerProvider(provider: designers.TableDesignerProvider): vscode.Disposable; export function registerExecutionPlanProvider(provider: executionPlan.ExecutionPlanProvider): vscode.Disposable; + /** + * Registers a server contextualization provider, which can provide context about a server to extensions like GitHub + * Copilot for improved suggestions. + * @param provider The provider to register + */ + export function registerServerContextualizationProvider(provider: contextualization.ServerContextualizationProvider): vscode.Disposable } export namespace designers { @@ -1773,6 +1780,29 @@ declare module 'azdata' { } } + export namespace contextualization { + export interface GetServerContextualizationResult { + /** + * An array containing the generated server context. + */ + context: string[]; + } + + export interface ServerContextualizationProvider extends DataProvider { + /** + * Generates server context. + * @param ownerUri The URI of the connection to generate context for. + */ + generateServerContextualization(ownerUri: string): void; + + /** + * Gets server context, which can be in the form of create scripts but is left up each provider. + * @param ownerUri The URI of the connection to get context for. + */ + getServerContextualization(ownerUri: string): Thenable; + } + } + /** * Component to display text with an icon representing the severity */ diff --git a/src/sql/platform/query/common/query.ts b/src/sql/platform/query/common/query.ts index 73f6fb8891..d70ccc957d 100644 --- a/src/sql/platform/query/common/query.ts +++ b/src/sql/platform/query/common/query.ts @@ -46,6 +46,7 @@ export interface IQueryEditorConfiguration { readonly tabColorMode: 'off' | 'border' | 'fill'; readonly showConnectionInfoInTitle: boolean; readonly promptToSaveGeneratedFiles: boolean; + readonly githubCopilotContextualizationEnabled: boolean; } export interface IResultGridConfiguration { diff --git a/src/sql/workbench/api/browser/mainThreadDataProtocol.ts b/src/sql/workbench/api/browser/mainThreadDataProtocol.ts index 2f2cf98e5a..d0de82ca4f 100644 --- a/src/sql/workbench/api/browser/mainThreadDataProtocol.ts +++ b/src/sql/workbench/api/browser/mainThreadDataProtocol.ts @@ -33,6 +33,7 @@ import { ITableDesignerService } from 'sql/workbench/services/tableDesigner/comm import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces'; import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers'; import { SqlExtHostContext, SqlMainContext } from 'vs/workbench/api/common/extHost.protocol'; +import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces'; /** * Main thread class for handling data protocol management registration. @@ -64,7 +65,8 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData @IDataGridProviderService private _dataGridProviderService: IDataGridProviderService, @IAdsTelemetryService private _telemetryService: IAdsTelemetryService, @ITableDesignerService private _tableDesignerService: ITableDesignerService, - @IExecutionPlanService private _executionPlanService: IExecutionPlanService + @IExecutionPlanService private _executionPlanService: IExecutionPlanService, + @IServerContextualizationService private _serverContextualizationService: IServerContextualizationService ) { super(); if (extHostContext) { @@ -571,6 +573,14 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData }); } + // Database server contextualization handler + public $registerServerContextualizationProvider(providerId: string, handle: number): void { + this._serverContextualizationService.registerProvider(providerId, { + generateServerContextualization: (ownerUri: string) => this._proxy.$generateServerContextualization(handle, ownerUri), + getServerContextualization: (ownerUri: string) => this._proxy.$getServerContextualization(handle, ownerUri) + }); + } + // Connection Management handlers public $onConnectionComplete(handle: number, connectionInfoSummary: azdata.ConnectionInfoSummary): void { this._connectionManagementService.onConnectionComplete(handle, connectionInfoSummary); diff --git a/src/sql/workbench/api/common/extHostDataProtocol.ts b/src/sql/workbench/api/common/extHostDataProtocol.ts index 2f22800b5c..98a13bf620 100644 --- a/src/sql/workbench/api/common/extHostDataProtocol.ts +++ b/src/sql/workbench/api/common/extHostDataProtocol.ts @@ -210,6 +210,12 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape { return rt; } + $registerServerContextualizationProvider(provider: azdata.contextualization.ServerContextualizationProvider): vscode.Disposable { + let rt = this.registerProvider(provider, DataProviderType.ServerContextualizationProvider); + this._proxy.$registerServerContextualizationProvider(provider.providerId, provider.handle); + return rt; + } + // Capabilities Discovery handlers override $getServerCapabilities(handle: number, client: azdata.DataProtocolClientCapabilities): Thenable { return this._resolveProvider(handle).getServerCapabilities(client); @@ -963,4 +969,14 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape { public override $isExecutionPlan(handle: number, value: string): Thenable { return this._resolveProvider(handle).isExecutionPlan(value); } + + // Database Server Contextualization API + + public override $generateServerContextualization(handle: number, ownerUri: string): void { + this._resolveProvider(handle).generateServerContextualization(ownerUri); + } + + public override $getServerContextualization(handle: number, ownerUri: string): Thenable { + return this._resolveProvider(handle).getServerContextualization(ownerUri); + } } diff --git a/src/sql/workbench/api/common/sqlExtHost.api.impl.ts b/src/sql/workbench/api/common/sqlExtHost.api.impl.ts index 083b516470..e88ee30d44 100644 --- a/src/sql/workbench/api/common/sqlExtHost.api.impl.ts +++ b/src/sql/workbench/api/common/sqlExtHost.api.impl.ts @@ -408,6 +408,10 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp return extHostDataProvider.$registerExecutionPlanProvider(provider); }; + let registerServerContextualizationProvider = (provider: azdata.contextualization.ServerContextualizationProvider): vscode.Disposable => { + return extHostDataProvider.$registerServerContextualizationProvider(provider); + }; + // namespace: dataprotocol const dataprotocol: typeof azdata.dataprotocol = { registerBackupProvider, @@ -430,6 +434,7 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp registerDataGridProvider, registerTableDesignerProvider, registerExecutionPlanProvider: registerExecutionPlanProvider, + registerServerContextualizationProvider: registerServerContextualizationProvider, onDidChangeLanguageFlavor(listener: (e: azdata.DidChangeLanguageFlavorParams) => any, thisArgs?: any, disposables?: extHostTypes.Disposable[]) { return extHostDataProvider.onDidChangeLanguageFlavor(listener, thisArgs, disposables); }, diff --git a/src/sql/workbench/api/common/sqlExtHost.protocol.ts b/src/sql/workbench/api/common/sqlExtHost.protocol.ts index 7bdeee8119..6e700cfeda 100644 --- a/src/sql/workbench/api/common/sqlExtHost.protocol.ts +++ b/src/sql/workbench/api/common/sqlExtHost.protocol.ts @@ -595,6 +595,14 @@ export abstract class ExtHostDataProtocolShape { * Determines if the provided value is an execution plan and returns the appropriate file extension. */ $isExecutionPlan(handle: number, value: string): Thenable { throw ni(); } + /** + * Generates server context. + */ + $generateServerContextualization(handle: number, ownerUri: string): void { throw ni(); } + /** + * Gets server context. + */ + $getServerContextualization(handle: number, ownerUri: string): Thenable { throw ni(); } } /** @@ -687,6 +695,7 @@ export interface MainThreadDataProtocolShape extends IDisposable { $registerDataGridProvider(providerId: string, title: string, handle: number): void; $registerTableDesignerProvider(providerId: string, handle: number): Promise; $registerExecutionPlanProvider(providerId: string, handle: number): void; + $registerServerContextualizationProvider(providerId: string, handle: number): void; $unregisterProvider(handle: number): Promise; $onConnectionComplete(handle: number, connectionInfoSummary: azdata.ConnectionInfoSummary): void; $onIntelliSenseCacheComplete(handle: number, connectionUri: string): void; diff --git a/src/sql/workbench/api/common/sqlExtHostTypes.ts b/src/sql/workbench/api/common/sqlExtHostTypes.ts index 8d745cf2d4..0186c9be61 100644 --- a/src/sql/workbench/api/common/sqlExtHostTypes.ts +++ b/src/sql/workbench/api/common/sqlExtHostTypes.ts @@ -421,7 +421,8 @@ export enum DataProviderType { SqlAssessmentServicesProvider = 'SqlAssessmentServicesProvider', DataGridProvider = 'DataGridProvider', TableDesignerProvider = 'TableDesignerProvider', - ExecutionPlanProvider = 'ExecutionPlanProvider' + ExecutionPlanProvider = 'ExecutionPlanProvider', + ServerContextualizationProvider = 'ServerContextualizationProvider' } export enum DeclarativeDataType { diff --git a/src/sql/workbench/browser/editor/query/fileQueryEditorInput.ts b/src/sql/workbench/browser/editor/query/fileQueryEditorInput.ts index 380a758c01..aeb05d5fe1 100644 --- a/src/sql/workbench/browser/editor/query/fileQueryEditorInput.ts +++ b/src/sql/workbench/browser/editor/query/fileQueryEditorInput.ts @@ -18,6 +18,8 @@ import { FILE_QUERY_EDITOR_TYPEID } from 'sql/workbench/common/constants'; import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation'; import { EditorInput } from 'vs/workbench/common/editor/editorInput'; import { IResourceEditorInput } from 'vs/platform/editor/common/editor'; +import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces'; +import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions'; export class FileQueryEditorInput extends QueryEditorInput { @@ -30,9 +32,11 @@ export class FileQueryEditorInput extends QueryEditorInput { @IConnectionManagementService connectionManagementService: IConnectionManagementService, @IQueryModelService queryModelService: IQueryModelService, @IConfigurationService configurationService: IConfigurationService, - @IInstantiationService instantiationService: IInstantiationService + @IInstantiationService instantiationService: IInstantiationService, + @IServerContextualizationService serverContextualizationService: IServerContextualizationService, + @IExtensionService extensionService: IExtensionService ) { - super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService); + super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService, extensionService); } public override resolve(): Promise { diff --git a/src/sql/workbench/browser/editor/query/untitledQueryEditorInput.ts b/src/sql/workbench/browser/editor/query/untitledQueryEditorInput.ts index ad0bfa39c0..3d6f802360 100644 --- a/src/sql/workbench/browser/editor/query/untitledQueryEditorInput.ts +++ b/src/sql/workbench/browser/editor/query/untitledQueryEditorInput.ts @@ -22,6 +22,8 @@ import { IUntitledQueryEditorInput } from 'sql/workbench/common/editor/query/unt import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService'; import { Uri } from 'vscode'; import { ILogService } from 'vs/platform/log/common/log'; +import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces'; +import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions'; export class UntitledQueryEditorInput extends QueryEditorInput implements IUntitledQueryEditorInput { @@ -36,9 +38,11 @@ export class UntitledQueryEditorInput extends QueryEditorInput implements IUntit @IConfigurationService configurationService: IConfigurationService, @IInstantiationService instantiationService: IInstantiationService, @ILogService private readonly logService: ILogService, - @IEditorResolverService private readonly editorResolverService: IEditorResolverService + @IEditorResolverService private readonly editorResolverService: IEditorResolverService, + @IServerContextualizationService serverContextualizationService: IServerContextualizationService, + @IExtensionService extensionService: IExtensionService ) { - super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService); + super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService, extensionService); // Set the mode explicitely to stop the auto language detection service from changing the mode unexpectedly. // the auto language detection service won't do the language change only if the mode is explicitely set. // if the mode (e.g. kusto, sql) do not exist for whatever reason, we will default it to sql. diff --git a/src/sql/workbench/common/editor/query/queryEditorInput.ts b/src/sql/workbench/common/editor/query/queryEditorInput.ts index f0a29a6315..eff674a0ec 100644 --- a/src/sql/workbench/common/editor/query/queryEditorInput.ts +++ b/src/sql/workbench/common/editor/query/queryEditorInput.ts @@ -20,6 +20,8 @@ import { AbstractTextResourceEditorInput } from 'vs/workbench/common/editor/text import { IQueryEditorConfiguration } from 'sql/platform/query/common/query'; import { EditorInput } from 'vs/workbench/common/editor/editorInput'; import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation'; +import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces'; +import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions'; const MAX_SIZE = 13; @@ -142,6 +144,8 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab private _state = this._register(new QueryEditorState()); public get state(): QueryEditorState { return this._state; } + private _serverContext: string[]; + constructor( private _description: string | undefined, protected _text: AbstractTextResourceEditorInput, @@ -149,7 +153,9 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab @IConnectionManagementService private readonly connectionManagementService: IConnectionManagementService, @IQueryModelService private readonly queryModelService: IQueryModelService, @IConfigurationService private readonly configurationService: IConfigurationService, - @IInstantiationService protected readonly instantiationService: IInstantiationService + @IInstantiationService protected readonly instantiationService: IInstantiationService, + @IServerContextualizationService private readonly serverContextualizationService: IServerContextualizationService, + @IExtensionService private readonly extensionService: IExtensionService ) { super(); @@ -235,6 +241,27 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab public override isDirty(): boolean { return this._text.isDirty(); } public get resource(): URI { return this._text.resource; } + public async getServerContext(): Promise { + const copilotExt = await this.extensionService.getExtension('github.copilot'); + + if (copilotExt && this.configurationService.getValue('queryEditor').githubCopilotContextualizationEnabled) { + if (!this._serverContext) { + const result = await this.serverContextualizationService.getServerContextualization(this.uri); + // TODO lewissanchez - Remove this from here once Copilot starts pulling context. That isn't implemented yet, so + // getting scripts this way for now. + this._serverContext = result.context; + + return this._serverContext; + } + else { + return this._serverContext; + } + } + else { + return Promise.resolve([]); + } + } + public override getName(longForm?: boolean): string { if (this.configurationService.getValue('queryEditor').showConnectionInfoInTitle) { let profile = this.connectionManagementService.getConnectionProfile(this.uri); diff --git a/src/sql/workbench/contrib/commandLine/test/electron-browser/commandLine.test.ts b/src/sql/workbench/contrib/commandLine/test/electron-browser/commandLine.test.ts index a07bd46bdd..cce62f977c 100644 --- a/src/sql/workbench/contrib/commandLine/test/electron-browser/commandLine.test.ts +++ b/src/sql/workbench/contrib/commandLine/test/electron-browser/commandLine.test.ts @@ -467,7 +467,7 @@ suite('commandLineService tests', () => { let uri = URI.file(args._[0]); const workbenchinstantiationService = workbenchInstantiationService(); const editorInput = workbenchinstantiationService.createInstance(FileEditorInput, uri, undefined, undefined, undefined, undefined, undefined, undefined); - const queryInput = new FileQueryEditorInput(undefined, editorInput, undefined, connectionManagementService.object, querymodelService.object, configurationService.object, workbenchinstantiationService); + const queryInput = new FileQueryEditorInput(undefined, editorInput, undefined, connectionManagementService.object, querymodelService.object, configurationService.object, workbenchinstantiationService, undefined, undefined); queryInput.state.connected = true; const editorService: TypeMoq.Mock = TypeMoq.Mock.ofType(TestEditorService, TypeMoq.MockBehavior.Strict); editorService.setup(e => e.editors).returns(() => [queryInput]); diff --git a/src/sql/workbench/contrib/query/browser/query.contribution.ts b/src/sql/workbench/contrib/query/browser/query.contribution.ts index 92b4a552d1..90ae87146c 100644 --- a/src/sql/workbench/contrib/query/browser/query.contribution.ts +++ b/src/sql/workbench/contrib/query/browser/query.contribution.ts @@ -374,6 +374,11 @@ const queryEditorConfiguration: IConfigurationNode = { 'type': 'boolean', 'default': false, 'description': localize('queryEditor.promptToSaveGeneratedFiles', "Prompt to save generated SQL files") + }, + 'queryEditor.githubCopilotContextualizationEnabled': { + 'type': 'boolean', + 'default': false, + 'description': localize('queryEditor.githubCopilotContextualizationEnabled', "(Preview) Enable contextualization of queries for GitHub Copilot. This setting helps GitHub Copilot to return improved suggestions, if the Copilot extension is installed and providers have implemented contextualization.") } } }; diff --git a/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts b/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts index 1078b2e7ac..56fbef5ceb 100644 --- a/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts +++ b/src/sql/workbench/contrib/query/test/browser/queryEditor.test.ts @@ -316,6 +316,8 @@ suite('SQL QueryEditor Tests', () => { configurationService.object, testinstantiationService, undefined, + undefined, + undefined, undefined ); }); diff --git a/src/sql/workbench/services/contextualization/common/interfaces.ts b/src/sql/workbench/services/contextualization/common/interfaces.ts new file mode 100644 index 0000000000..b53dc15f14 --- /dev/null +++ b/src/sql/workbench/services/contextualization/common/interfaces.ts @@ -0,0 +1,43 @@ +/*--------------------------------------------------------------------------------------------- + * 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 { createDecorator } from 'vs/platform/instantiation/common/instantiation'; + + +export const SERVICE_ID = 'serverContextualizationService'; +export const IServerContextualizationService = createDecorator(SERVICE_ID); + +export interface IServerContextualizationService { + _serviceBrand: undefined; + + /** + * Register a server contextualization service provider + */ + registerProvider(providerId: string, provider: azdata.contextualization.ServerContextualizationProvider): void; + + /** + * Unregister a server contextualization service provider + */ + unregisterProvider(providerId: string): void; + + /** + * 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 + */ + getProvider(providerId: string): azdata.contextualization.ServerContextualizationProvider; + + /** + * Generates server context + * @param ownerUri The URI of the connection to generate context for. + */ + generateServerContextualization(ownerUri: string): void; + + /** + * Gets all database context. + * @param ownerUri The URI of the connection to get context for. + */ + getServerContextualization(ownerUri: string): Promise; +} diff --git a/src/sql/workbench/services/contextualization/common/serverContextualizationService.ts b/src/sql/workbench/services/contextualization/common/serverContextualizationService.ts new file mode 100644 index 0000000000..a9980d4b75 --- /dev/null +++ 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(); + + 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('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 { + const providerName = this._connectionManagementService.getProviderIdFromUri(ownerUri); + const handler = this.getProvider(providerName); + if (handler) { + return await handler.getServerContextualization(ownerUri); + } + else { + return Promise.resolve({ + context: [] + }); + } + } +} diff --git a/src/vs/workbench/workbench.common.main.ts b/src/vs/workbench/workbench.common.main.ts index 01d8c85d86..57d3be4ccc 100644 --- a/src/vs/workbench/workbench.common.main.ts +++ b/src/vs/workbench/workbench.common.main.ts @@ -228,6 +228,8 @@ import { ITableDesignerService } from 'sql/workbench/services/tableDesigner/comm import { TableDesignerService } from 'sql/workbench/services/tableDesigner/browser/tableDesignerService'; import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces'; import { ExecutionPlanService } from 'sql/workbench/services/executionPlan/common/executionPlanService'; +import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces'; +import { ServerContextualizationService } from 'sql/workbench/services/contextualization/common/serverContextualizationService'; import { IErrorDiagnosticsService } from 'sql/workbench/services/diagnostics/common/errorDiagnosticsService'; import { ErrorDiagnosticsService } from 'sql/workbench/services/diagnostics/browser/errorDiagnosticsService'; import { IComponentContextService, ComponentContextService } from 'sql/workbench/services/componentContext/browser/componentContextService'; @@ -275,6 +277,7 @@ registerSingleton(IAssessmentService, AssessmentService, InstantiationType.Eager registerSingleton(IDataGridProviderService, DataGridProviderService, InstantiationType.Eager); registerSingleton(ITableDesignerService, TableDesignerService, InstantiationType.Eager); registerSingleton(IExecutionPlanService, ExecutionPlanService, InstantiationType.Eager); +registerSingleton(IServerContextualizationService, ServerContextualizationService, InstantiationType.Eager); registerSingleton(IErrorDiagnosticsService, ErrorDiagnosticsService, InstantiationType.Eager); registerSingleton(IComponentContextService, ComponentContextService, InstantiationType.Eager); //#endregion