/*--------------------------------------------------------------------------------------------- * 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: [] }); } } }