Add acquireLoginSession API for Azdata (#13985)

* wip

* fix tests

* add tests

* PR comments

extensions/arc/src/models/controllerModel.ts

@@ -61,10 +61,11 @@ export class ControllerModel {
 	}
 
 	/**
-	 * Calls azdata login to set the context to this controller
+	 * Calls azdata login to set the context to this controller and acquires a login session to prevent other
+	 * calls from changing the context while commands for this session are being executed.
 	 * @param promptReconnect
 	 */
-	public async azdataLogin(promptReconnect: boolean = false): Promise<void> {
+	public async acquireAzdataSession(promptReconnect: boolean = false): Promise<azdataExt.AzdataSession> {
 		let promptForValidClusterContext: boolean = false;
 		try {
 			const contexts = await getKubeConfigClusterContexts(this.info.kubeConfigFilePath);
@@ -102,7 +103,7 @@ export class ControllerModel {
 			}
 		}
 
-		await this._azdataApi.azdata.login(this.info.url, this.info.username, this._password, this.azdataAdditionalEnvVars);
+		return this._azdataApi.azdata.acquireSession(this.info.url, this.info.username, this._password, this.azdataAdditionalEnvVars);
 	}
 
 	/**
@@ -123,64 +124,68 @@ export class ControllerModel {
 		}
 		// create a new in progress promise object
 		ControllerModel._refreshInProgress = new Deferred<void>();
-		await this.azdataLogin(promptReconnect);
+		const session = await this.acquireAzdataSession(promptReconnect);
 		const newRegistrations: Registration[] = [];
-		await Promise.all([
-			this._azdataApi.azdata.arc.dc.config.show().then(result => {
-				this._controllerConfig = result.result;
-				this.configLastUpdated = new Date();
-				this._onConfigUpdated.fire(this._controllerConfig);
-			}).catch(err => {
-				// If an error occurs show a message so the user knows something failed but still
-				// fire the event so callers can know to update (e.g. so dashboards don't show the
-				// loading icon forever)
-				if (showErrors) {
-					vscode.window.showErrorMessage(loc.fetchConfigFailed(this.info.name, err));
-				}
-				this._onConfigUpdated.fire(this._controllerConfig);
-				throw err;
-			}),
-			this._azdataApi.azdata.arc.dc.endpoint.list(this.azdataAdditionalEnvVars).then(result => {
-				this._endpoints = result.result;
-				this.endpointsLastUpdated = new Date();
-				this._onEndpointsUpdated.fire(this._endpoints);
-			}).catch(err => {
-				// If an error occurs show a message so the user knows something failed but still
-				// fire the event so callers can know to update (e.g. so dashboards don't show the
-				// loading icon forever)
-				if (showErrors) {
-					vscode.window.showErrorMessage(loc.fetchEndpointsFailed(this.info.name, err));
-				}
-				this._onEndpointsUpdated.fire(this._endpoints);
-				throw err;
-			}),
-			Promise.all([
-				this._azdataApi.azdata.arc.postgres.server.list(this.azdataAdditionalEnvVars).then(result => {
-					newRegistrations.push(...result.result.map(r => {
-						return {
-							instanceName: r.name,
-							state: r.state,
-							instanceType: ResourceType.postgresInstances
-						};
-					}));
+		try {
+			await Promise.all([
+				this._azdataApi.azdata.arc.dc.config.show(this.azdataAdditionalEnvVars, session).then(result => {
+					this._controllerConfig = result.result;
+					this.configLastUpdated = new Date();
+					this._onConfigUpdated.fire(this._controllerConfig);
+				}).catch(err => {
+					// If an error occurs show a message so the user knows something failed but still
+					// fire the event so callers hooking into this can handle the error (e.g. so dashboards don't show the
+					// loading icon forever)
+					if (showErrors) {
+						vscode.window.showErrorMessage(loc.fetchConfigFailed(this.info.name, err));
+					}
+					this._onConfigUpdated.fire(this._controllerConfig);
+					throw err;
 				}),
-				this._azdataApi.azdata.arc.sql.mi.list().then(result => {
-					newRegistrations.push(...result.result.map(r => {
-						return {
-							instanceName: r.name,
-							state: r.state,
-							instanceType: ResourceType.sqlManagedInstances
-						};
-					}));
+				this._azdataApi.azdata.arc.dc.endpoint.list(this.azdataAdditionalEnvVars, session).then(result => {
+					this._endpoints = result.result;
+					this.endpointsLastUpdated = new Date();
+					this._onEndpointsUpdated.fire(this._endpoints);
+				}).catch(err => {
+					// If an error occurs show a message so the user knows something failed but still
+					// fire the event so callers can know to update (e.g. so dashboards don't show the
+					// loading icon forever)
+					if (showErrors) {
+						vscode.window.showErrorMessage(loc.fetchEndpointsFailed(this.info.name, err));
+					}
+					this._onEndpointsUpdated.fire(this._endpoints);
+					throw err;
+				}),
+				Promise.all([
+					this._azdataApi.azdata.arc.postgres.server.list(this.azdataAdditionalEnvVars, session).then(result => {
+						newRegistrations.push(...result.result.map(r => {
+							return {
+								instanceName: r.name,
+								state: r.state,
+								instanceType: ResourceType.postgresInstances
+							};
+						}));
+					}),
+					this._azdataApi.azdata.arc.sql.mi.list(this.azdataAdditionalEnvVars, session).then(result => {
+						newRegistrations.push(...result.result.map(r => {
+							return {
+								instanceName: r.name,
+								state: r.state,
+								instanceType: ResourceType.sqlManagedInstances
+							};
+						}));
+					})
+				]).then(() => {
+					this._registrations = newRegistrations;
+					this.registrationsLastUpdated = new Date();
+					this._onRegistrationsUpdated.fire(this._registrations);
 				})
-			]).then(() => {
-				this._registrations = newRegistrations;
-				this.registrationsLastUpdated = new Date();
-				this._onRegistrationsUpdated.fire(this._registrations);
-			})
-		]);
-		ControllerModel._refreshInProgress.resolve();
-		ControllerModel._refreshInProgress = undefined;
+			]);
+		} finally {
+			session.dispose();
+			ControllerModel._refreshInProgress.resolve();
+			ControllerModel._refreshInProgress = undefined;
+		}
 	}
 
 	public get endpoints(): azdataExt.DcEndpointListResult[] {

extensions/arc/src/models/miaaModel.ts

@@ -71,10 +71,11 @@ export class MiaaModel extends ResourceModel {
 			return this._refreshPromise.promise;
 		}
 		this._refreshPromise = new Deferred();
+		let session: azdataExt.AzdataSession | undefined = undefined;
 		try {
-			await this.controllerModel.azdataLogin();
+			session = await this.controllerModel.acquireAzdataSession();
 			try {
-				const result = await this._azdataApi.azdata.arc.sql.mi.show(this.info.name);
+				const result = await this._azdataApi.azdata.arc.sql.mi.show(this.info.name, this.controllerModel.azdataAdditionalEnvVars, session);
 				this._config = result.result;
 				this.configLastUpdated = new Date();
 				this._onConfigUpdated.fire(this._config);
@@ -114,6 +115,7 @@ export class MiaaModel extends ResourceModel {
 			this._refreshPromise.reject(err);
 			throw err;
 		} finally {
+			session?.dispose();
 			this._refreshPromise = undefined;
 		}
 	}

extensions/arc/src/models/postgresModel.ts

@@ -107,10 +107,10 @@ export class PostgresModel extends ResourceModel {
 			return this._refreshPromise.promise;
 		}
 		this._refreshPromise = new Deferred();
-
+		let session: azdataExt.AzdataSession | undefined = undefined;
 		try {
-			await this.controllerModel.azdataLogin();
-			this._config = (await this._azdataApi.azdata.arc.postgres.server.show(this.info.name)).result;
+			session = await this.controllerModel.acquireAzdataSession();
+			this._config = (await this._azdataApi.azdata.arc.postgres.server.show(this.info.name, this.controllerModel.azdataAdditionalEnvVars, session)).result;
 			this.configLastUpdated = new Date();
 			this._onConfigUpdated.fire(this._config);
 			this._refreshPromise.resolve();
@@ -118,6 +118,7 @@ export class PostgresModel extends ResourceModel {
 			this._refreshPromise.reject(err);
 			throw err;
 		} finally {
+			session?.dispose();
 			this._refreshPromise = undefined;
 		}
 	}

extensions/arc/src/test/mocks/fakeAzdataApi.ts

@@ -75,9 +75,12 @@ export class FakeAzdataApi implements azdataExt.IAzdataApi {
 	getPath(): Promise<string> {
 		throw new Error('Method not implemented.');
 	}
-	login(_endpoint: string, _username: string, _password: string): Promise<azdataExt.AzdataOutput<any>> {
+	login(_endpoint: string, _username: string, _password: string): Promise<azdataExt.AzdataOutput<void>> {
 		return <any>undefined;
 	}
+	acquireSession(_endpoint: string, _username: string, _password: string): Promise<azdataExt.AzdataSession> {
+		return Promise.resolve({ dispose: () => { } });
+	}
 	version(): Promise<azdataExt.AzdataOutput<string>> {
 		throw new Error('Method not implemented.');
 	}

extensions/arc/src/test/models/controllerModel.test.ts

@@ -43,7 +43,7 @@ describe('ControllerModel', function (): void {
 			// Returning an undefined model here indicates that the dialog closed without clicking "Ok" - usually through the user clicking "Cancel"
 			sinon.stub(ConnectToControllerDialog.prototype, 'waitForClose').returns(Promise.resolve(undefined));
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', kubeConfigFilePath: '/path/to/.kube/config', kubeClusterContext: 'currentCluster', username: 'admin', name: 'arc', rememberPassword: true, resources: [] });
-			await should(model.azdataLogin()).be.rejectedWith(new UserCancelledError(loc.userCancelledError));
+			await should(model.acquireAzdataSession()).be.rejectedWith(new UserCancelledError(loc.userCancelledError));
 		});
 
 		it('Reads password from cred store', async function (): Promise<void> {
@@ -58,13 +58,13 @@ describe('ControllerModel', function (): void {
 
 			const azdataExtApiMock = TypeMoq.Mock.ofType<azdataExt.IExtension>();
 			const azdataMock = TypeMoq.Mock.ofType<azdataExt.IAzdataApi>();
-			azdataMock.setup(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
+			azdataMock.setup(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
 			azdataExtApiMock.setup(x => x.azdata).returns(() => azdataMock.object);
 			sinon.stub(vscode.extensions, 'getExtension').returns(<any>{ exports: azdataExtApiMock.object });
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', kubeConfigFilePath: '/path/to/.kube/config', kubeClusterContext: 'currentCluster', username: 'admin', name: 'arc', rememberPassword: true, resources: [] });
 
-			await model.azdataLogin();
-			azdataMock.verify(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
+			await model.acquireAzdataSession();
+			azdataMock.verify(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 
 		it('Prompt for password when not in cred store', async function (): Promise<void> {
@@ -79,7 +79,7 @@ describe('ControllerModel', function (): void {
 
 			const azdataExtApiMock = TypeMoq.Mock.ofType<azdataExt.IExtension>();
 			const azdataMock = TypeMoq.Mock.ofType<azdataExt.IAzdataApi>();
-			azdataMock.setup(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
+			azdataMock.setup(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
 			azdataExtApiMock.setup(x => x.azdata).returns(() => azdataMock.object);
 			sinon.stub(vscode.extensions, 'getExtension').returns(<any>{ exports: azdataExtApiMock.object });
 
@@ -89,8 +89,8 @@ describe('ControllerModel', function (): void {
 
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', kubeConfigFilePath: '/path/to/.kube/config', kubeClusterContext: 'currentCluster', username: 'admin', name: 'arc', rememberPassword: true, resources: [] });
 
-			await model.azdataLogin();
-			azdataMock.verify(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
+			await model.acquireAzdataSession();
+			azdataMock.verify(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 
 		it('Prompt for password when rememberPassword is true but prompt reconnect is true', async function (): Promise<void> {
@@ -104,7 +104,7 @@ describe('ControllerModel', function (): void {
 
 			const azdataExtApiMock = TypeMoq.Mock.ofType<azdataExt.IExtension>();
 			const azdataMock = TypeMoq.Mock.ofType<azdataExt.IAzdataApi>();
-			azdataMock.setup(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
+			azdataMock.setup(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
 			azdataExtApiMock.setup(x => x.azdata).returns(() => azdataMock.object);
 			sinon.stub(vscode.extensions, 'getExtension').returns(<any>{ exports: azdataExtApiMock.object });
 
@@ -114,9 +114,9 @@ describe('ControllerModel', function (): void {
 
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', kubeConfigFilePath: '/path/to/.kube/config', kubeClusterContext: 'currentCluster', username: 'admin', name: 'arc', rememberPassword: true, resources: [] });
 
-			await model.azdataLogin(true);
+			await model.acquireAzdataSession(true);
 			should(waitForCloseStub.called).be.true('waitForClose should have been called');
-			azdataMock.verify(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
+			azdataMock.verify(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 
 		it('Prompt for password when we already have a password but prompt reconnect is true', async function (): Promise<void> {
@@ -130,7 +130,7 @@ describe('ControllerModel', function (): void {
 
 			const azdataExtApiMock = TypeMoq.Mock.ofType<azdataExt.IExtension>();
 			const azdataMock = TypeMoq.Mock.ofType<azdataExt.IAzdataApi>();
-			azdataMock.setup(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
+			azdataMock.setup(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
 			azdataExtApiMock.setup(x => x.azdata).returns(() => azdataMock.object);
 			sinon.stub(vscode.extensions, 'getExtension').returns(<any>{ exports: azdataExtApiMock.object });
 
@@ -141,9 +141,9 @@ describe('ControllerModel', function (): void {
 			// Set up original model with a password
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', kubeConfigFilePath: '/path/to/.kube/config', kubeClusterContext: 'currentCluster', username: 'admin', name: 'arc', rememberPassword: true, resources: [] }, 'originalPassword');
 
-			await model.azdataLogin(true);
+			await model.acquireAzdataSession(true);
 			should(waitForCloseStub.called).be.true('waitForClose should have been called');
-			azdataMock.verify(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
+			azdataMock.verify(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), password, TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 
 		it('Model values are updated correctly when modified during reconnect', async function (): Promise<void> {
@@ -158,7 +158,7 @@ describe('ControllerModel', function (): void {
 
 			const azdataExtApiMock = TypeMoq.Mock.ofType<azdataExt.IExtension>();
 			const azdataMock = TypeMoq.Mock.ofType<azdataExt.IAzdataApi>();
-			azdataMock.setup(x => x.login(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
+			azdataMock.setup(x => x.acquireSession(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => <any>Promise.resolve(undefined));
 			azdataExtApiMock.setup(x => x.azdata).returns(() => azdataMock.object);
 			sinon.stub(vscode.extensions, 'getExtension').returns(<any>{ exports: azdataExtApiMock.object });
 
@@ -199,10 +199,11 @@ describe('ControllerModel', function (): void {
 			const waitForCloseStub = sinon.stub(ConnectToControllerDialog.prototype, 'waitForClose').returns(Promise.resolve(
 				{ controllerModel: newModel, password: newPassword }));
 
-			await model.azdataLogin(true);
+			await model.acquireAzdataSession(true);
 			should(waitForCloseStub.called).be.true('waitForClose should have been called');
 			should((await treeDataProvider.getChildren()).length).equal(1, 'Tree Data provider should still only have 1 node');
 			should(model.info).deepEqual(newInfo, 'Model info should have been updated');
+
 		});
 
 	});

extensions/arc/src/ui/dashboards/miaa/miaaComputeAndStoragePage.ts

@@ -129,13 +129,16 @@ export class MiaaComputeAndStoragePage extends DashboardPage {
 							cancellable: false
 						},
 						async (_progress, _token): Promise<void> => {
+							let session: azdataExt.AzdataSession | undefined = undefined;
 							try {
-								await this._miaaModel.controllerModel.azdataLogin();
+								session = await this._miaaModel.controllerModel.acquireAzdataSession();
 								await this._azdataApi.azdata.arc.sql.mi.edit(
-									this._miaaModel.info.name, this.saveArgs);
+									this._miaaModel.info.name, this.saveArgs, this._miaaModel.controllerModel.azdataAdditionalEnvVars, session);
 							} catch (err) {
 								this.saveButton!.enabled = true;
 								throw err;
+							} finally {
+								session?.dispose();
 							}
 
 							await this._miaaModel.refresh();

extensions/arc/src/ui/dashboards/miaa/miaaDashboardOverviewPage.ts

@@ -206,8 +206,13 @@ export class MiaaDashboardOverviewPage extends DashboardPage {
 								cancellable: false
 							},
 							async (_progress, _token) => {
-								await this._controllerModel.azdataLogin();
-								return await this._azdataApi.azdata.arc.sql.mi.delete(this._miaaModel.info.name);
+								const session = await this._controllerModel.acquireAzdataSession();
+								try {
+									return await this._azdataApi.azdata.arc.sql.mi.delete(this._miaaModel.info.name, this._controllerModel.azdataAdditionalEnvVars, session);
+								} finally {
+									session.dispose();
+								}
+
 							}
 						);
 						await this._controllerModel.refreshTreeNode();

extensions/arc/src/ui/dashboards/postgres/postgresComputeAndStoragePage.ts

@@ -155,18 +155,23 @@ export class PostgresComputeAndStoragePage extends DashboardPage {
 							cancellable: false
 						},
 						async (_progress, _token): Promise<void> => {
+							let session: azdataExt.AzdataSession | undefined = undefined;
 							try {
-								await this._postgresModel.controllerModel.azdataLogin();
+								session = await this._postgresModel.controllerModel.acquireAzdataSession();
 								await this._azdataApi.azdata.arc.postgres.server.edit(
 									this._postgresModel.info.name,
 									this.saveArgs,
-									this._postgresModel.engineVersion
+									this._postgresModel.engineVersion,
+									this._postgresModel.controllerModel.azdataAdditionalEnvVars,
+									session
 								);
 							} catch (err) {
 								// If an error occurs while editing the instance then re-enable the save button since
 								// the edit wasn't successfully applied
 								this.saveButton!.enabled = true;
 								throw err;
+							} finally {
+								session?.dispose();
 							}
 							await this._postgresModel.refresh();
 						}

extensions/arc/src/ui/dashboards/postgres/postgresOverviewPage.ts

@@ -151,16 +151,21 @@ export class PostgresOverviewPage extends DashboardPage {
 				try {
 					const password = await promptAndConfirmPassword(input => !input ? loc.enterANonEmptyPassword : '');
 					if (password) {
-						await this._postgresModel.controllerModel.azdataLogin();
-						await this._azdataApi.azdata.arc.postgres.server.edit(
-							this._postgresModel.info.name,
-							{
-								adminPassword: true,
-								noWait: true
-							},
-							this._postgresModel.engineVersion,
-							{ 'AZDATA_PASSWORD': password }
-						);
+						const session = await this._postgresModel.controllerModel.acquireAzdataSession();
+						try {
+							await this._azdataApi.azdata.arc.postgres.server.edit(
+								this._postgresModel.info.name,
+								{
+									adminPassword: true,
+									noWait: true
+								},
+								this._postgresModel.engineVersion,
+								Object.assign({ 'AZDATA_PASSWORD': password }, this._controllerModel.azdataAdditionalEnvVars),
+								session
+							);
+						} finally {
+							session.dispose();
+						}
 						vscode.window.showInformationMessage(loc.passwordReset);
 					}
 				} catch (error) {
@@ -188,8 +193,13 @@ export class PostgresOverviewPage extends DashboardPage {
 								cancellable: false
 							},
 							async (_progress, _token) => {
-								await this._postgresModel.controllerModel.azdataLogin();
-								return await this._azdataApi.azdata.arc.postgres.server.delete(this._postgresModel.info.name);
+								const session = await this._postgresModel.controllerModel.acquireAzdataSession();
+								try {
+									return await this._azdataApi.azdata.arc.postgres.server.delete(this._postgresModel.info.name, this._controllerModel.azdataAdditionalEnvVars, session);
+								} finally {
+									session.dispose();
+								}
+
 							}
 						);
 						await this._controllerModel.refreshTreeNode();

extensions/arc/src/ui/dashboards/postgres/postgresParametersPage.ts

@@ -172,8 +172,14 @@ export class PostgresParametersPage extends DashboardPage {
 								this.engineSettingUpdates!.forEach((value: string) => {
 									this.engineSettings += value + ', ';
 								});
-								await this._azdataApi.azdata.arc.postgres.server.edit(
-									this._postgresModel.info.name, { engineSettings: this.engineSettings + `'` });
+								const session = await this._postgresModel.controllerModel.acquireAzdataSession();
+								try {
+									await this._azdataApi.azdata.arc.postgres.server.edit(
+										this._postgresModel.info.name, { engineSettings: this.engineSettings + `'` });
+								} finally {
+									session.dispose();
+								}
+
 							} catch (err) {
 								// If an error occurs while editing the instance then re-enable the save button since
 								// the edit wasn't successfully applied
@@ -237,7 +243,9 @@ export class PostgresParametersPage extends DashboardPage {
 						async (_progress, _token): Promise<void> => {
 							//all
 							// azdata arc postgres server edit -n <server group name> -e '' -re
+							let session: azdataExt.AzdataSession | undefined = undefined;
 							try {
+								session = await this._postgresModel.controllerModel.acquireAzdataSession();
 								await this._azdataApi.azdata.arc.postgres.server.edit(
 									this._postgresModel.info.name, { engineSettings: `'' -re` });
 							} catch (err) {
@@ -245,6 +253,8 @@ export class PostgresParametersPage extends DashboardPage {
 								// the edit wasn't successfully applied
 								this.resetButton!.enabled = true;
 								throw err;
+							} finally {
+								session?.dispose();
 							}
 							await this._postgresModel.refresh();
 						}
@@ -463,9 +473,14 @@ export class PostgresParametersPage extends DashboardPage {
 							title: loc.updatingInstance(this._postgresModel.info.name),
 							cancellable: false
 						},
-						(_progress, _token) => {
-							return this._azdataApi.azdata.arc.postgres.server.edit(
-								this._postgresModel.info.name, { engineSettings: name + '=' });
+						async (_progress, _token) => {
+							const session = await this._postgresModel.controllerModel.acquireAzdataSession();
+							try {
+								this._azdataApi.azdata.arc.postgres.server.edit(
+									this._postgresModel.info.name, { engineSettings: name + '=' });
+							} finally {
+								session.dispose();
+							}
 						}
 					);