Fix query retry bug (#17066)

* fixed retry query bug

* lint fixes

* pr fixes

* fix package.json

* pr fixes

* fix package.json

* pr fixes

* fix cannot connect message

localization/xliff/enu/constants/localizedConstants.enu.xlf

@@ -125,6 +125,9 @@
         <trans-unit id="azureAddAccount">
             <source xml:lang="en">Add an Account...</source>
         </trans-unit>
+        <trans-unit id="cannotConnect">
+            <source xml:lang="en">Cannot connect due to expired tokens. Please re-authenticate and try again.</source>
+        </trans-unit>
         <trans-unit id="aad">
             <source xml:lang="en">AAD</source>
         </trans-unit>

package.json

@@ -106,8 +106,8 @@
     "typescript": "^3.5.3",
     "uglify-js": "mishoo/UglifyJS2#harmony-v2.8.22",
     "vsce": "^1.58.0",
-    "vscode-test": "^1.0.0",
     "vscode-nls-dev": "2.0.1",
+    "vscode-test": "^1.0.0",
     "xmldom": "^0.6.0",
     "yargs": "https://registry.npmjs.org/yargs/-/yargs-3.32.0.tgz"
   },

src/controllers/connectionManager.ts

@@ -29,6 +29,7 @@ import { QuestionTypes, IQuestion } from '../prompts/question';
 import { IAccount } from '../models/contracts/azure/accountInterfaces';
 import { AzureController } from '../azure/azureController';
 import { IConnectionInfo } from 'vscode-mssql';
+import providerSettings from '../azure/providerSettings';
 
 /**
  * Information for a document's connection. Exported for testing purposes.
@@ -672,6 +673,27 @@ export default class ConnectionManager {
     // create a new connection with the connectionCreds provided
     public async connect(fileUri: string, connectionCreds: IConnectionInfo, promise?: Deferred<boolean>): Promise<boolean> {
         const self = this;
+        // Check if the azure account token is present before sending connect request
+        if (connectionCreds.authenticationType === Constants.azureMfa) {
+            if (!connectionCreds.azureAccountToken) {
+                let account = this.accountStore.getAccount(connectionCreds.accountId);
+                let profile = new ConnectionProfile(connectionCreds);
+                let azureAccountToken = await this.azureController.refreshToken(account, this.accountStore, providerSettings.resources.databaseResource);
+                if (!azureAccountToken) {
+                    let errorMessage = LocalizedConstants.msgAccountRefreshFailed;
+                    let refreshResult = await this.vscodeWrapper.showErrorMessage(errorMessage, LocalizedConstants.refreshTokenLabel);
+                    if (refreshResult === LocalizedConstants.refreshTokenLabel) {
+                        await this.azureController.getTokens(
+                            profile, this.accountStore, providerSettings.resources.databaseResource);
+
+                    } else {
+                        throw new Error(`${LocalizedConstants.cannotConnect}`);
+                    }
+                } else {
+                    connectionCreds.azureAccountToken = azureAccountToken;
+                }
+            }
+        }
         let connectionPromise = new Promise<boolean>(async (resolve, reject) => {
             let connectionInfo: ConnectionInfo = new ConnectionInfo();
             connectionInfo.credentials = connectionCreds;

src/models/connectionProfile.ts

@@ -30,6 +30,18 @@ export class ConnectionProfile extends ConnectionCredentials implements IConnect
     public accountStore: AccountStore;
     public accountId: string;
 
+    constructor(connectionCredentials?: ConnectionCredentials) {
+        super();
+        if (connectionCredentials) {
+            this.accountId = connectionCredentials.accountId;
+            this.authenticationType = connectionCredentials.authenticationType;
+            this.azureAccountToken = connectionCredentials.azureAccountToken;
+            this.database = connectionCredentials.database;
+            this.email = connectionCredentials.email;
+            this.password = connectionCredentials.password;
+            this.server = connectionCredentials.server;
+        }
+    }
     /**
      * Creates a new profile by prompting the user for information.
      * @param  {IPrompter} prompter that asks user the questions needed to complete a profile

src/objectExplorer/objectExplorerService.ts

@@ -420,14 +420,7 @@ export class ObjectExplorerService {
                 } else if (connectionCredentials.authenticationType === Constants.azureMfa) {
                     let azureController = this._connectionManager.azureController;
                     let account = this._connectionManager.accountStore.getAccount(connectionCredentials.accountId);
-                    let profile = new ConnectionProfile();
-                    profile.accountId = connectionCredentials.accountId;
-                    profile.authenticationType = connectionCredentials.authenticationType;
-                    profile.azureAccountToken = connectionCredentials.azureAccountToken;
-                    profile.database = connectionCredentials.database;
-                    profile.email = connectionCredentials.email;
-                    profile.password = connectionCredentials.password;
-                    profile.server = connectionCredentials.server;
+                    let profile = new ConnectionProfile(connectionCredentials);
                     if (!connectionCredentials.azureAccountToken) {
                         let azureAccountToken = await azureController.refreshToken(
                             account, this._connectionManager.accountStore, providerSettings.resources.databaseResource);