API improvement: make registerConnectionEventProvider return disposable (#11880)

* promote api to official

* add comments

* disposable

* move getConnection out

* comment for connection namespace

* remove extra line

* also fix registerQueryInfoHandler

src/sql/azdata.d.ts

@@ -4072,7 +4072,7 @@ declare module 'azdata' {
 		/**
 		 * Register a query event listener
 		 */
-		export function registerQueryEventListener(listener: QueryEventListener): void;
+		export function registerQueryEventListener(listener: QueryEventListener): vscode.Disposable;
 
 		/**
 		 * Get a QueryDocument object for a file URI

src/sql/azdata.proposed.d.ts

@@ -13,21 +13,39 @@ declare module 'azdata' {
 	 * Namespace for connection management
 	 */
 	export namespace connection {
+		/**
+		 * Supported connection event types
+		 */
 		export type ConnectionEventType =
 			| 'onConnect'
 			| 'onDisconnect'
 			| 'onConnectionChanged';
 
+		/**
+		 * Connection Event Lister
+		 */
 		export interface ConnectionEventListener {
+			/**
+			 * Connection event handler
+			 * @param type Connection event type
+			 * @param ownerUri Connection's owner uri
+			 * @param args Connection profile
+			 */
 			onConnectionEvent(type: ConnectionEventType, ownerUri: string, args: IConnectionProfile): void;
 		}
 
 		/**
 		 * Register a connection event listener
+		 * @param listener The connection event listener
 		 */
-		export function registerConnectionEventListener(listener: connection.ConnectionEventListener): void;
+		export function registerConnectionEventListener(listener: connection.ConnectionEventListener): vscode.Disposable;
 
-		export function getConnection(uri: string): Thenable<ConnectionProfile>;
+		/**
+		 * Get connection profile by its owner uri
+		 * @param ownerUri The owner uri of the connection
+		 * @returns Promise to return the connection profile matching the ownerUri
+		 */
+		export function getConnection(ownerUri: string): Thenable<ConnectionProfile>;
 	}
 
 	export namespace nb {

src/sql/workbench/api/browser/mainThreadConnectionManagement.ts

@@ -12,7 +12,7 @@ import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/br
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import * as TaskUtilities from 'sql/workbench/browser/taskUtilities';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { isUndefinedOrNull } from 'vs/base/common/types';
 import { generateUuid } from 'vs/base/common/uuid';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
@@ -24,6 +24,7 @@ import { deepClone } from 'vs/base/common/objects';
 export class MainThreadConnectionManagement extends Disposable implements MainThreadConnectionManagementShape {
 
 	private _proxy: ExtHostConnectionManagementShape;
+	private _connectionEventListenerDisposables = new Map<number, IDisposable>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -39,7 +40,7 @@ export class MainThreadConnectionManagement extends Disposable implements MainTh
 		}
 	}
 
-	public $registerConnectionEventListener(handle: number, providerId: string): void {
+	public $registerConnectionEventListener(handle: number): void {
 
 		let stripProfile = (inputProfile: azdata.IConnectionProfile) => {
 			if (!inputProfile) {
@@ -66,17 +67,28 @@ export class MainThreadConnectionManagement extends Disposable implements MainTh
 			return outputProfile;
 		};
 
-		this._connectionManagementService.onConnect((params: IConnectionParams) => {
+		const disposable = new DisposableStore();
+		disposable.add(this._connectionManagementService.onConnect((params: IConnectionParams) => {
 			this._proxy.$onConnectionEvent(handle, 'onConnect', params.connectionUri, stripProfile(params.connectionProfile));
-		});
+		}));
 
-		this._connectionManagementService.onConnectionChanged((params: IConnectionParams) => {
+		disposable.add(this._connectionManagementService.onConnectionChanged((params: IConnectionParams) => {
 			this._proxy.$onConnectionEvent(handle, 'onConnectionChanged', params.connectionUri, stripProfile(params.connectionProfile));
-		});
+		}));
 
-		this._connectionManagementService.onDisconnect((params: IConnectionParams) => {
+		disposable.add(this._connectionManagementService.onDisconnect((params: IConnectionParams) => {
 			this._proxy.$onConnectionEvent(handle, 'onDisconnect', params.connectionUri, stripProfile(params.connectionProfile));
-		});
+		}));
+
+		this._connectionEventListenerDisposables.set(handle, disposable);
+	}
+
+	public $unregisterConnectionEventListener(handle: number): void {
+		const disposable = this._connectionEventListenerDisposables.get(handle);
+		if (disposable) {
+			disposable.dispose();
+			this._connectionEventListenerDisposables.delete(handle);
+		}
 	}
 
 	public $getConnections(activeConnectionsOnly?: boolean): Thenable<azdata.connection.ConnectionProfile[]> {

src/sql/workbench/api/browser/mainThreadQueryEditor.ts

@@ -9,7 +9,7 @@ import { extHostNamedCustomer } from 'vs/workbench/api/common/extHostCustomers';
 import { IConnectionManagementService, IConnectionCompletionOptions, ConnectionType, RunQueryOnConnectionMode } from 'sql/platform/connection/common/connectionManagement';
 import { QueryEditor } from 'sql/workbench/contrib/query/browser/queryEditor';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 import { IQueryModelService } from 'sql/workbench/services/query/common/queryModel';
 import * as azdata from 'azdata';
 import { IQueryManagementService } from 'sql/workbench/services/query/common/queryManagement';
@@ -21,6 +21,7 @@ import { ILogService } from 'vs/platform/log/common/log';
 export class MainThreadQueryEditor extends Disposable implements MainThreadQueryEditorShape {
 
 	private _proxy: ExtHostQueryEditorShape;
+	private _queryEventListenerDisposables = new Map<number, IDisposable>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -112,10 +113,19 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 	}
 
 	public $registerQueryInfoListener(handle: number): void {
-		this._register(this._queryModelService.onQueryEvent(event => {
+		const disposable = this._queryModelService.onQueryEvent(event => {
 			let connectionProfile = this._connectionManagementService.getConnectionProfile(event.uri);
 			this._proxy.$onQueryEvent(connectionProfile?.providerName, handle, event.uri, event);
-		}));
+		});
+		this._queryEventListenerDisposables.set(handle, disposable);
+	}
+
+	public $unregisterQueryInfoListener(handle: number): void {
+		const disposable = this._queryEventListenerDisposables.get(handle);
+		if (disposable) {
+			disposable.dispose();
+			this._queryEventListenerDisposables.delete(handle);
+		}
 	}
 
 	public $createQueryTab(fileUri: string, title: string, componentId: string): void {

src/sql/workbench/api/common/extHostConnectionManagement.ts

@@ -6,6 +6,8 @@
 import { ExtHostConnectionManagementShape, SqlMainContext, MainThreadConnectionManagementShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
 import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
 import * as azdata from 'azdata';
+import { IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable } from 'vs/workbench/api/common/extHostTypes';
 
 export class ExtHostConnectionManagement extends ExtHostConnectionManagementShape {
 
@@ -27,10 +29,15 @@ export class ExtHostConnectionManagement extends ExtHostConnectionManagementShap
 		}
 	}
 
-	public $registerConnectionEventListener(providerId: string, listener: azdata.connection.ConnectionEventListener): void {
+	public $registerConnectionEventListener(listener: azdata.connection.ConnectionEventListener): IDisposable {
-		this._connectionListeners[this._nextListenerHandle] = listener;
+		const handle = this._nextListenerHandle++;
-		this._proxy.$registerConnectionEventListener(this._nextListenerHandle, providerId);
+		this._connectionListeners[handle] = listener;
-		this._nextListenerHandle++;
+		this._proxy.$registerConnectionEventListener(handle);
+
+		return new Disposable(() => {
+			this._connectionListeners.delete(handle);
+			this._proxy.$unregisterConnectionEventListener(handle);
+		});
 	}
 
 	public $getCurrentConnection(): Thenable<azdata.connection.ConnectionProfile> {

src/sql/workbench/api/common/extHostQueryEditor.ts

@@ -8,6 +8,7 @@ import { ExtHostQueryEditorShape, SqlMainContext, MainThreadQueryEditorShape } f
 import * as azdata from 'azdata';
 import { IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
+import { Disposable } from 'vs/workbench/api/common/extHostTypes';
 
 class ExtHostQueryDocument implements azdata.queryeditor.QueryDocument {
 	constructor(
@@ -52,10 +53,14 @@ export class ExtHostQueryEditor implements ExtHostQueryEditorShape {
 		return this._proxy.$runQuery(fileUri, runCurrentQuery);
 	}
 
-	public $registerQueryInfoListener(listener: azdata.queryeditor.QueryEventListener): void {
+	public $registerQueryInfoListener(listener: azdata.queryeditor.QueryEventListener): Disposable {
-		this._queryListeners[this._nextListenerHandle] = listener;
+		const handle = this._nextListenerHandle++;
-		this._proxy.$registerQueryInfoListener(this._nextListenerHandle);
+		this._queryListeners[handle] = listener;
-		this._nextListenerHandle++;
+		this._proxy.$registerQueryInfoListener(handle);
+		return new Disposable(() => {
+			this._queryListeners.delete(handle);
+			this._proxy.$unregisterQueryInfoListener(handle);
+		});
 	}
 
 	public $onQueryEvent(providerId: string, handle: number, fileUri: string, event: IQueryEvent): void {

src/sql/workbench/api/common/sqlExtHost.api.impl.ts

@@ -27,7 +27,6 @@ import { ExtHostNotebookDocumentsAndEditors } from 'sql/workbench/api/common/ext
 import { ExtHostExtensionManagement } from 'sql/workbench/api/common/extHostExtensionManagement';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
-import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
 import { IURITransformerService } from 'vs/workbench/api/common/extHostUriTransformerService';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
@@ -107,8 +106,8 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
 				getConnections(activeConnectionsOnly?: boolean): Thenable<azdata.connection.ConnectionProfile[]> {
 					return extHostConnectionManagement.$getConnections(activeConnectionsOnly);
 				},
-				registerConnectionEventListener(listener: azdata.connection.ConnectionEventListener): void {
+				registerConnectionEventListener(listener: azdata.connection.ConnectionEventListener): vscode.Disposable {
-					return extHostConnectionManagement.$registerConnectionEventListener(mssqlProviderName, listener);
+					return extHostConnectionManagement.$registerConnectionEventListener(listener);
 				},
 				getConnection(uri: string): Thenable<azdata.connection.ConnectionProfile> {
 					return extHostConnectionManagement.$getConnection(uri);
@@ -485,8 +484,8 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
 					extHostQueryEditor.$runQuery(fileUri, runCurrentQuery);
 				},
 
-				registerQueryEventListener(listener: azdata.queryeditor.QueryEventListener): void {
+				registerQueryEventListener(listener: azdata.queryeditor.QueryEventListener): extHostTypes.Disposable {
-					extHostQueryEditor.$registerQueryInfoListener(listener);
+					return extHostQueryEditor.$registerQueryInfoListener(listener);
 				},
 
 				getQueryDocument(fileUri: string): Thenable<azdata.queryeditor.QueryDocument> {

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -615,7 +615,8 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 }
 
 export interface MainThreadConnectionManagementShape extends IDisposable {
-	$registerConnectionEventListener(handle: number, providerId: string): void;
+	$registerConnectionEventListener(handle: number): void;
+	$unregisterConnectionEventListener(handle: number): void;
 	$getConnections(activeConnectionsOnly?: boolean): Thenable<azdata.connection.ConnectionProfile[]>;
 	$getConnection(uri: string): Thenable<azdata.connection.ConnectionProfile>;
 	$getActiveConnections(): Thenable<azdata.connection.Connection[]>;
@@ -830,6 +831,7 @@ export interface MainThreadQueryEditorShape extends IDisposable {
 	$createQueryTab(fileUri: string, title: string, content: string): void;
 	$setQueryExecutionOptions(fileUri: string, options: azdata.QueryExecutionOptions): Thenable<void>;
 	$registerQueryInfoListener(handle: number): void;
+	$unregisterQueryInfoListener(handle: number): void;
 }
 
 export interface ExtHostNotebookShape {