From 0e4c67dc0c7fe92c44a97c23d2329b79c792c8e9 Mon Sep 17 00:00:00 2001 From: Hai Cao Date: Tue, 21 Sep 2021 18:57:54 -0700 Subject: [PATCH] Add AAD token expiration handling for query runner (#17117) * Add AAD token refresh for notebook * move token refresh to query management and remove previous refresh calls --- src/sql/workbench/contrib/query/browser/queryActions.ts | 3 --- src/sql/workbench/services/query/common/queryManagement.ts | 4 +++- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/sql/workbench/contrib/query/browser/queryActions.ts b/src/sql/workbench/contrib/query/browser/queryActions.ts index 1773aad273..f429f06707 100644 --- a/src/sql/workbench/contrib/query/browser/queryActions.ts +++ b/src/sql/workbench/contrib/query/browser/queryActions.ts @@ -233,7 +233,6 @@ export class RunQueryAction extends QueryTaskbarAction { editor = this.editor; } - await this.connectionManagementService.refreshAzureAccountTokenIfNecessary(this.editor.input.uri); if (this.isConnected(editor)) { // Hide IntelliSense suggestions list when running query to match SSMS behavior this.commandService?.executeCommand('hideSuggestWidget'); @@ -306,7 +305,6 @@ export class EstimatedQueryPlanAction extends QueryTaskbarAction { public override async run(): Promise { if (!this.editor.isSelectionEmpty()) { - await this.connectionManagementService.refreshAzureAccountTokenIfNecessary(this.editor.input.uri); if (this.isConnected(this.editor)) { // If we are already connected, run the query this.runQuery(this.editor); @@ -346,7 +344,6 @@ export class ActualQueryPlanAction extends QueryTaskbarAction { public override async run(): Promise { if (!this.editor.isSelectionEmpty()) { - await this.connectionManagementService.refreshAzureAccountTokenIfNecessary(this.editor.input.uri); if (this.isConnected(this.editor)) { // If we are already connected, run the query this.runQuery(this.editor); diff --git a/src/sql/workbench/services/query/common/queryManagement.ts b/src/sql/workbench/services/query/common/queryManagement.ts index 5b7e662cc3..a7b9e80256 100644 --- a/src/sql/workbench/services/query/common/queryManagement.ts +++ b/src/sql/workbench/services/query/common/queryManagement.ts @@ -212,7 +212,9 @@ export class QueryManagementService implements IQueryManagementService { } let handler = this._requestHandlers.get(providerId); if (handler) { - return action(handler); + return this._connectionService.refreshAzureAccountTokenIfNecessary(uri).then(() => { + return action(handler); + }); } else { return Promise.reject(new Error('No Handler Registered')); }