error handling (#12140)

extensions/azurecore/src/azureResource/utils.ts

@@ -140,7 +140,11 @@ export async function getResourceGroups(appContext: AppContext, account?: azdata
 	return result;
 }
 
-export async function runResourceQuery<T extends azureResource.AzureGraphResource>(appContext: AppContext, account: azdata.Account, subscription: azureResource.AzureResourceSubscription, ignoreErrors: boolean = false, query: string) {
+export async function runResourceQuery<T extends azureResource.AzureGraphResource>(
+	account: azdata.Account,
+	subscriptions: azureResource.AzureResourceSubscription[],
+	ignoreErrors: boolean = false,
+	query: string): Promise<ResourceQueryResult<T>> {
 	const result: ResourceQueryResult<T> = { resources: [], errors: [] };
 	if (!account?.properties?.tenants || !isArray(account.properties.tenants)) {
 		const error = new Error(localize('azure.accounts.runResourceQuery.errors.invalidAccount', "Invalid account"));
@@ -151,51 +155,74 @@ export async function runResourceQuery<T extends azureResource.AzureGraphResourc
 		return result;
 	}
 
-	if (!subscription.tenant) {
+	// Check our subscriptions to ensure we have valid ones
-		const error = new Error(localize('azure.accounts.runResourceQuery.errors.noTenantSpecifiedForSubscription', "Invalid tenant for subscription"));
+	subscriptions.forEach(subscription => {
-		if (!ignoreErrors) {
+		if (!subscription.tenant) {
-			throw error;
+			const error = new Error(localize('azure.accounts.runResourceQuery.errors.noTenantSpecifiedForSubscription', "Invalid tenant for subscription"));
+			if (!ignoreErrors) {
+				throw error;
+			}
+			result.errors.push(error);
 		}
-		result.errors.push(error);
+	});
+	if (result.errors.length > 0) {
 		return result;
 	}
 
-	const tokenResponse = await azdata.accounts.getAccountSecurityToken(account, subscription.tenant, azdata.AzureResource.ResourceManagement);
+	// We need to get a different security token for each tenant to query the resources for the subscriptions on
-	const token = tokenResponse.token;
+	// that tenant
-	const tokenType = tokenResponse.tokenType;
+	for (let i = 0; i < account.properties.tenants.length; ++i) {
-	const credential = new TokenCredentials(token, tokenType);
+		const tenant = account.properties.tenants[i];
-
+		const tenantSubscriptions = subscriptions.filter(subscription => subscription.tenant === tenant.id);
-	const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
+		if (tenantSubscriptions.length < 1) {
-
+			// We may not have all subscriptions or the tenant might not have any subscriptions - just ignore these ones
-	const allResources: T[] = [];
+			continue;
-	let totalProcessed = 0;
-
-	const doQuery = async (skipToken?: string) => {
-		const response = await resourceClient.resources({
-			subscriptions: [subscription.id],
-			query,
-			options: {
-				resultFormat: 'objectArray',
-				skipToken: skipToken
-			}
-		});
-		const resources: T[] = response.data;
-		totalProcessed += resources.length;
-		allResources.push(...resources);
-		if (response.skipToken && totalProcessed < response.totalRecords) {
-			await doQuery(response.skipToken);
 		}
-	};
-	try {
-		await doQuery();
-	} catch (err) {
-		console.error(err);
-		const error = new Error(localize('azure.accounts.runResourceQuery.errors.invalidQuery', "Invalid query"));
-		result.errors.push(error);
-	}
-	result.resources = allResources;
-	return result;
 
+		let resourceClient: ResourceGraphClient;
+		try {
+			const tokenResponse = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.ResourceManagement);
+			const token = tokenResponse.token;
+			const tokenType = tokenResponse.tokenType;
+			const credential = new TokenCredentials(token, tokenType);
+
+			resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
+		} catch (err) {
+			console.error(err);
+			const error = new Error(localize('azure.accounts.runResourceQuery.errors.unableToFetchToken', "Unable to get token for tenant {0}", tenant.id));
+			result.errors.push(error);
+			continue;
+		}
+
+		const allResources: T[] = [];
+		let totalProcessed = 0;
+
+		const doQuery = async (skipToken?: string) => {
+			const response = await resourceClient.resources({
+				subscriptions: tenantSubscriptions.map(subscription => subscription.id),
+				query,
+				options: {
+					resultFormat: 'objectArray',
+					skipToken: skipToken
+				}
+			});
+			const resources: T[] = response.data;
+			totalProcessed += resources.length;
+			allResources.push(...resources);
+			if (response.skipToken && totalProcessed < response.totalRecords) {
+				await doQuery(response.skipToken);
+			}
+		};
+		try {
+			await doQuery();
+		} catch (err) {
+			console.error(err);
+			const error = new Error(localize('azure.accounts.runResourceQuery.errors.invalidQuery', "Invalid query"));
+			result.errors.push(error);
+		}
+		result.resources.push(...allResources);
+	}
+	return result;
 }
 
 export async function getSubscriptions(appContext: AppContext, account?: azdata.Account, ignoreErrors: boolean = false): Promise<GetSubscriptionsResult> {

extensions/azurecore/src/azurecore.d.ts

@@ -72,7 +72,7 @@ declare module 'azurecore' {
 		getRegionDisplayName(region?: string): string;
 		provideResources(): azureResource.IAzureResourceProvider[];
 
-		runGraphQuery<T extends azureResource.AzureGraphResource>(account: azdata.Account, subscription: azureResource.AzureResourceSubscription, ignoreErrors: boolean, query: string): Promise<ResourceQueryResult<T>>;
+		runGraphQuery<T extends azureResource.AzureGraphResource>(account: azdata.Account, subscriptions: azureResource.AzureResourceSubscription[], ignoreErrors: boolean, query: string): Promise<ResourceQueryResult<T>>;
 	}
 
 	export type GetSubscriptionsResult = { subscriptions: azureResource.AzureResourceSubscription[], errors: Error[] };

extensions/azurecore/src/extension.ts

@@ -110,10 +110,10 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 		},
 		getRegionDisplayName: utils.getRegionDisplayName,
 		runGraphQuery<T extends azureResource.AzureGraphResource>(account: azdata.Account,
-			subscription: azureResource.AzureResourceSubscription,
+			subscriptions: azureResource.AzureResourceSubscription[],
 			ignoreErrors: boolean,
 			query: string): Promise<azurecore.ResourceQueryResult<T>> {
-			return azureResourceUtils.runResourceQuery(appContext, account, subscription, ignoreErrors, query);
+			return azureResourceUtils.runResourceQuery(account, subscriptions, ignoreErrors, query);
 		}
 	};
 }

extensions/machine-learning/src/test/stubs.ts

@@ -8,7 +8,7 @@ import * as azurecore from 'azurecore';
 import { azureResource } from 'azureResource';
 
 export class AzurecoreApiStub implements azurecore.IExtension {
-	runGraphQuery<T extends azureResource.AzureGraphResource>(_account: azdata.Account, _subscription: azureResource.AzureResourceSubscription, _ignoreErrors: boolean, _query: string): Promise<azurecore.ResourceQueryResult<T>> {
+	runGraphQuery<T extends azureResource.AzureGraphResource>(_account: azdata.Account, _subscriptions: azureResource.AzureResourceSubscription[], _ignoreErrors: boolean, _query: string): Promise<azurecore.ResourceQueryResult<T>> {
 		throw new Error('Method not implemented.');
 	}
 	getSubscriptions(_account?: azdata.Account | undefined, _ignoreErrors?: boolean | undefined): Thenable<azurecore.GetSubscriptionsResult> {

extensions/sql-migration/src/api/azure.ts

@@ -39,7 +39,7 @@ export type SqlManagedInstance = AzureProduct;
 export async function getAvailableManagedInstanceProducts(account: azdata.Account, subscription: Subscription): Promise<SqlManagedInstance[]> {
 	const api = await getAzureCoreAPI();
 
-	const result = await api.runGraphQuery<SqlManagedInstance>(account, subscription, false, 'where type == "microsoft.sql/managedinstances"');
+	const result = await api.runGraphQuery<SqlManagedInstance>(account, [subscription], false, 'where type == "microsoft.sql/managedinstances"');
 	return result.resources;
 }
 
@@ -47,7 +47,7 @@ export type SqlServer = AzureProduct;
 export async function getAvailableSqlServers(account: azdata.Account, subscription: Subscription): Promise<SqlServer[]> {
 	const api = await getAzureCoreAPI();
 
-	const result = await api.runGraphQuery<SqlServer>(account, subscription, false, 'where type == "microsoft.sql/servers"');
+	const result = await api.runGraphQuery<SqlServer>(account, [subscription], false, 'where type == "microsoft.sql/servers"');
 	return result.resources;
 }
 
@@ -55,6 +55,6 @@ export type SqlVMServer = AzureProduct;
 export async function getAvailableSqlVMs(account: azdata.Account, subscription: Subscription): Promise<SqlVMServer[]> {
 	const api = await getAzureCoreAPI();
 
-	const result = await api.runGraphQuery<SqlVMServer>(account, subscription, false, 'where type == "microsoft.compute/virtualmachines" and properties.storageProfile.imageReference.publisher == "microsoftsqlserver"');
+	const result = await api.runGraphQuery<SqlVMServer>(account, [subscription], false, 'where type == "microsoft.compute/virtualmachines" and properties.storageProfile.imageReference.publisher == "microsoftsqlserver"');
 	return result.resources;
 }