3644 Kusto Token Refresh (#12576)

* 3644 Added RequestSecurityTokenParams, RequestSecurityTokenResponse, and SecurityTokenRequest to Kusto/contracts.ts. Added AccountFeature to features.ts. Registered feature in kustoServer.ts

* 3644 Removed TryCatch in kusto features > getToken

* 3644 Added AccountId to Kusto > RequestSecurityTokenParams. Refactored kusto features getToken to use the accountId for the query window.

* 3644 Removed unused AccountQuickPickItem

extensions/kusto/src/contracts.ts

@@ -29,6 +29,24 @@ export class TelemetryParams {
 
 // ------------------------------- </ Telemetry Sent Event > ----------------------------------
 
+// ------------------------------- < Security Token Request > ------------------------------------------
+export interface RequestSecurityTokenParams {
+	authority: string;
+	provider: string;
+	resource: string;
+	accountId: string;
+}
+
+export interface RequestSecurityTokenResponse {
+	accountKey: string;
+	token: string;
+}
+
+export namespace SecurityTokenRequest {
+	export const type = new RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse, void, void>('account/securityTokenRequest');
+}
+// ------------------------------- </ Security Token Request > ------------------------------------------
+
 // ------------------------------- <Serialization> -----------------------------
 export namespace SerializeDataStartRequest {
 	export const type = new RequestType<azdata.SerializeDataStartRequestParams, azdata.SerializeDataResult, void, void>('serialize/start');

extensions/kusto/src/features.ts

@@ -5,12 +5,13 @@
 
 import { SqlOpsDataClient, SqlOpsFeature } from 'dataprotocol-client';
 import { ClientCapabilities, StaticFeature, RPCMessageType, ServerCapabilities } from 'vscode-languageclient';
-import { Disposable } from 'vscode';
+import { Disposable, window } from 'vscode';
 import { Telemetry } from './telemetry';
 import * as contracts from './contracts';
 import * as azdata from 'azdata';
 import * as Utils from './utils';
 import * as UUID from 'vscode-languageclient/lib/utils/uuid';
+import { localize } from './localize';
 
 export class TelemetryFeature implements StaticFeature {
 
@@ -27,6 +28,57 @@ export class TelemetryFeature implements StaticFeature {
 	}
 }
 
+export class AccountFeature implements StaticFeature {
+
+	constructor(private _client: SqlOpsDataClient) { }
+
+	fillClientCapabilities(capabilities: ClientCapabilities): void { }
+
+	initialize(): void {
+		this._client.onRequest(contracts.SecurityTokenRequest.type, async (request): Promise<contracts.RequestSecurityTokenResponse | undefined> => {
+			return this.getToken(request);
+		});
+	}
+
+	protected async getToken(request: contracts.RequestSecurityTokenParams): Promise<contracts.RequestSecurityTokenResponse | undefined> {
+		const accountList = await azdata.accounts.getAllAccounts();
+		let account: azdata.Account | undefined;
+
+		if (accountList.length < 1) {
+			// TODO: Prompt user to add account
+			window.showErrorMessage(localize('kusto.missingLinkedAzureAccount', "Azure Data Studio needs to contact Azure Key Vault to access a column master key for Always Encrypted, but no linked Azure account is available. Please add a linked Azure account and retry the query."));
+			return undefined;
+		} else {
+			account = accountList.find(a => a.key.accountId === request.accountId);
+		}
+
+		if (!account) {
+			window.showErrorMessage(localize('kusto.accountDoesNotExist', "Account does not exist."));
+			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;
+		}
+		const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.Sql);
+
+		if (!securityToken?.token) {
+			window.showErrorMessage(unauthorizedMessage);
+			return undefined;
+		}
+
+		let params: contracts.RequestSecurityTokenResponse = {
+			accountKey: JSON.stringify(account.key),
+			token: securityToken.token
+		};
+
+		return params;
+	}
+}
+
 export class SerializationFeature extends SqlOpsFeature<undefined> {
 	private static readonly messageTypes: RPCMessageType[] = [
 		contracts.SerializeDataStartRequest.type,

extensions/kusto/src/kustoServer.ts

@@ -12,7 +12,7 @@ import { getCommonLaunchArgsAndCleanupOldLogFiles } from './utils';
 import { localize } from './localize';
 import { Telemetry, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
-import { TelemetryFeature, SerializationFeature } from './features';
+import { TelemetryFeature, SerializationFeature, AccountFeature } from './features';
 import { AppContext } from './appContext';
 import { CompletionExtensionParams, CompletionExtLoadRequest } from './contracts';
 import { promises as fs } from 'fs';
@@ -135,6 +135,7 @@ function getClientOptions(context: AppContext): ClientOptions {
 			// we only want to add new features
 			...SqlOpsDataClient.defaultFeatures,
 			TelemetryFeature,
+			AccountFeature,
 			SerializationFeature
 		],
 		outputChannel: new CustomOutputChannel()