Add support for clearing pooled connections (#24325)

extensions/mssql/package.json

@@ -157,6 +157,10 @@
           "dark": "resources/dark/groupBySchemaDisabled_inverse.svg",
           "light": "resources/light/groupBySchemaDisabled.svg"
         }
+      },
+      {
+        "command": "mssql.clearPooledConnections",
+        "title": "%mssql.connection.clearPooledConnections%"
       }
     ],
     "outputChannels": [

extensions/mssql/package.nls.json

@@ -174,9 +174,10 @@
 	"title.newTable": "New Table",
 	"title.designTable": "Design",
 	"title.changeNotebookConnection": "Change SQL Notebook Connection",
+	"mssql.connection.clearPooledConnections": "SQL Server: Clear Pooled Connections",
 	"mssql.parallelMessageProcessing": "[Experimental] Whether the requests to the SQL Tools Service should be handled in parallel. This is introduced to discover the issues there might be when handling all requests in parallel. The default value is false. Azure Data Studio is required to be relaunched when the value is changed.",
 	"mssql.enableSqlAuthenticationProvider": "Enables use of the Sql Authentication Provider for 'Active Directory Interactive' authentication mode when user selects 'AzureMFA' authentication. This enables Server-side resource endpoint integration when fetching access tokens. This option is only supported for 'MSAL' Azure Authentication Library. Azure Data Studio is required to be relaunched when the value is changed.",
-	"mssql.enableConnectionPooling": "Enables connection pooling on MSSQL connections to improve overall performance of Azure Data Studio connectivity. This setting is enabled by default. Azure Data Studio is required to be relaunched when the value is changed.",
+	"mssql.enableConnectionPooling": "Enables connection pooling on MSSQL connections to improve overall performance of Azure Data Studio connectivity. This setting is enabled by default. Azure Data Studio is required to be relaunched when the value is changed. To clear pooled connections, run the command: 'SQL Server: Clear Pooled Connections'",
 	"mssql.tableDesigner.preloadDatabaseModel": "Whether to preload the database model when the database node in the object explorer is expanded. When enabled, the loading time of table designer can be reduced. Note: You might see higher than normal memory usage if you need to expand a lot of database nodes.",
 	"mssql.tableDesigner.allowDisableAndReenableDdlTriggers": "Whether to allow table designer to disable and re-enable DDL triggers during publish",
 	"mssql.objectExplorer.groupBySchema": "When enabled, the database objects in Object Explorer will be categorized by schema.",

extensions/mssql/src/connection/commands.ts

@@ -0,0 +1,19 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { AppContext } from '../appContext';
+import * as vscode from 'vscode';
+import * as constants from './constants';
+import { ConnectionService } from './connectionService';
+
+export function registerConnectionCommands(appContext: AppContext) {
+	appContext.extensionContext.subscriptions.push(vscode.commands.registerCommand('mssql.clearPooledConnections', async () => {
+		await getConnectionService(appContext).clearPooledConnections();
+	}));
+}
+
+function getConnectionService(appContext: AppContext): ConnectionService {
+	return appContext.getService<ConnectionService>(constants.ConnectionService);
+}

extensions/mssql/src/connection/connectionService.ts

@@ -0,0 +1,33 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as constants from './constants';
+import * as contracts from '../contracts';
+
+import { BaseService, ISqlOpsFeature, SqlOpsDataClient } from 'dataprotocol-client';
+import { ClientCapabilities } from 'vscode-languageclient';
+import { AppContext } from '../appContext';
+
+export class ConnectionService extends BaseService {
+	public static asFeature(context: AppContext): ISqlOpsFeature {
+		return class extends ConnectionService {
+			constructor(client: SqlOpsDataClient) {
+				super(context, client);
+			}
+
+			fillClientCapabilities(_: ClientCapabilities): void { }
+			initialize(): void { }
+		};
+	}
+
+	private constructor(context: AppContext, client: SqlOpsDataClient) {
+		super(client);
+		context.registerService(constants.ConnectionService, this);
+	}
+
+	async clearPooledConnections(): Promise<void> {
+		return this.runWithErrorHandling(contracts.ClearPooledConnectionsRequest.type, {});
+	}
+}

extensions/mssql/src/connection/constants.ts

@@ -0,0 +1,6 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export const ConnectionService = 'ConnectionService';

extensions/mssql/src/contracts.ts

@@ -1733,6 +1733,11 @@ export namespace EncryptionKeysChangedNotification {
 	export const type = new NotificationType<DidChangeEncryptionIVKeyParams, void>('connection/encryptionKeysChanged');
 }
 
+// ------------------------------- < Clear Pooled Connections Request > ---------------------------------------
+
+export namespace ClearPooledConnectionsRequest {
+	export const type = new RequestType<object, void, void, void>('connection/clearpooledconnections');
+}
 // ------------------------------- < Query Store > ------------------------------------
 //#region Query Store
 

extensions/mssql/src/main.ts

@@ -27,6 +27,7 @@ import { registerTableDesignerCommands } from './tableDesigner/tableDesigner';
 import { registerObjectManagementCommands } from './objectManagement/commands';
 import { TelemetryActions, TelemetryReporter, TelemetryViews } from './telemetry';
 import { noConvertResult, noDocumentFound, unsupportedPlatform } from './localizedConstants';
+import { registerConnectionCommands } from './connection/commands';
 
 const localize = nls.loadMessageBundle();
 
@@ -156,6 +157,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<IExten
 
 	registerTableDesignerCommands(appContext);
 	registerObjectManagementCommands(appContext);
+	registerConnectionCommands(appContext);
 
 	// context.subscriptions.push(new SqlNotebookController()); Temporarily disabled due to breaking query editor
 

extensions/mssql/src/sqlToolsServer.ts

@@ -32,6 +32,7 @@ import { ErrorDiagnosticsProvider } from './errorDiagnostics/errorDiagnosticsPro
 import { SqlProjectsService } from './sqlProjects/sqlProjectsService';
 import { ObjectManagementService } from './objectManagement/objectManagementService';
 import { QueryStoreService } from './queryStore/queryStoreService';
+import { ConnectionService } from './connection/connectionService';
 
 const localize = nls.loadMessageBundle();
 const outputChannel = vscode.window.createOutputChannel(Constants.serviceName);
@@ -237,6 +238,7 @@ function getClientOptions(context: AppContext): ClientOptions {
 			AgentServicesFeature,
 			SerializationFeature,
 			SqlAssessmentServicesFeature,
+			ConnectionService.asFeature(context),
 			SchemaCompareService.asFeature(context),
 			LanguageExtensionService.asFeature(context),
 			DacFxService.asFeature(context),