From 7dfafe1393e420a7592a57c00cccf343c1a2c10c Mon Sep 17 00:00:00 2001
From: Justin M <63619224+JustinMDotNet@users.noreply.github.com>
Date: Tue, 30 Mar 2021 10:30:25 -0700
Subject: [PATCH] dstsAuth Token Refresh (#14890)

* Refactored getToken to set tenantId based on providerId

* Changed logic to set tenantId when provider is not dstsAuth
---
 extensions/kusto/src/features.ts | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/extensions/kusto/src/features.ts b/extensions/kusto/src/features.ts
index ffbb2420f2..be46e5d776 100644
--- a/extensions/kusto/src/features.ts
+++ b/extensions/kusto/src/features.ts
@@ -59,13 +59,19 @@ export class AccountFeature implements StaticFeature {
 			return undefined;
 		}
 
-		const tenant = account.properties.tenants.find((t: { [key: string]: string }) => request.authority.includes(t.id));
 		const unauthorizedMessage = localize('kusto.insufficientlyPrivelagedAzureAccount', "The configured Azure account for {0} does not have sufficient permissions for Azure Key Vault to access a column master key for Always Encrypted.", account.key.accountId);
-		if (!tenant) {
-			window.showErrorMessage(unauthorizedMessage);
-			return undefined;
+
+		let tenantId: string = '';
+		if (request.provider !== 'dstsAuth') {
+			const tenant = account.properties.tenants.find((t: { [key: string]: string }) => request.authority.includes(t.id));
+			if (!tenant) {
+				window.showErrorMessage(unauthorizedMessage);
+				return undefined;
+			}
+			tenantId = tenant.id;
 		}
-		const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.Sql);
+
+		const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenantId, azdata.AzureResource.Sql);
 
 		if (!securityToken?.token) {
 			window.showErrorMessage(unauthorizedMessage);