From 9004769865481bc06eff864a4cbdf7340a4b7f5e Mon Sep 17 00:00:00 2001 From: Charles Gagnon Date: Thu, 3 Sep 2020 19:07:37 -0700 Subject: [PATCH] Update azure graph queries to allow multiple subs (#12124) --- extensions/azurecore/src/azureResource/interfaces.ts | 2 +- .../providers/database/databaseService.ts | 8 ++++---- .../providers/resourceTreeDataProviderBase.ts | 10 +++++----- extensions/azurecore/src/azureResource/utils.ts | 2 +- .../database/databaseTreeDataProvider.test.ts | 2 +- .../databaseServerTreeDataProvider.test.ts | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/extensions/azurecore/src/azureResource/interfaces.ts b/extensions/azurecore/src/azureResource/interfaces.ts index 1e21f1e9e5..dc1163cf21 100644 --- a/extensions/azurecore/src/azureResource/interfaces.ts +++ b/extensions/azurecore/src/azureResource/interfaces.ts @@ -38,5 +38,5 @@ export interface IAzureResourceNodeWithProviderId { } export interface IAzureResourceService { - getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials, account: Account): Promise; + getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: msRest.ServiceClientCredentials, account: Account): Promise; } diff --git a/extensions/azurecore/src/azureResource/providers/database/databaseService.ts b/extensions/azurecore/src/azureResource/providers/database/databaseService.ts index f8ca86a622..6501faca90 100644 --- a/extensions/azurecore/src/azureResource/providers/database/databaseService.ts +++ b/extensions/azurecore/src/azureResource/providers/database/databaseService.ts @@ -15,13 +15,13 @@ interface DatabaseGraphData extends GraphData { kind: string; } export class AzureResourceDatabaseService implements IAzureResourceService { - public async getResources(subscription: azureResource.AzureResourceSubscription, credential: ServiceClientCredentials, account: Account): Promise { + public async getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: ServiceClientCredentials, account: Account): Promise { const databases: azureResource.AzureResourceDatabase[] = []; const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint }); // Query servers and databases in parallel (start both promises before waiting on the 1st) - let serverQueryPromise = queryGraphResources(resourceClient, subscription.id, serversQuery); - let dbQueryPromise = queryGraphResources(resourceClient, subscription.id, 'where type == "microsoft.sql/servers/databases"'); + let serverQueryPromise = queryGraphResources(resourceClient, subscriptions, serversQuery); + let dbQueryPromise = queryGraphResources(resourceClient, subscriptions, 'where type == "microsoft.sql/servers/databases"'); let servers: DbServerGraphData[] = await serverQueryPromise as DbServerGraphData[]; let dbByGraph: DatabaseGraphData[] = await dbQueryPromise as DatabaseGraphData[]; @@ -35,7 +35,7 @@ export class AzureResourceDatabaseService implements IAzureResourceService { // Filter master DBs, and for all others find their server to get login info diff --git a/extensions/azurecore/src/azureResource/providers/resourceTreeDataProviderBase.ts b/extensions/azurecore/src/azureResource/providers/resourceTreeDataProviderBase.ts index 34d32da86e..08e514911d 100644 --- a/extensions/azurecore/src/azureResource/providers/resourceTreeDataProviderBase.ts +++ b/extensions/azurecore/src/azureResource/providers/resourceTreeDataProviderBase.ts @@ -44,7 +44,7 @@ export abstract class ResourceTreeDataProviderBase[]; + const resources: T[] = await this._resourceService.getResources([element.subscription], credential, element.account) || []; return resources; } @@ -63,12 +63,12 @@ export interface GraphData { } -export async function queryGraphResources(resourceClient: ResourceGraphClient, subId: string, resourceQuery: string): Promise { +export async function queryGraphResources(resourceClient: ResourceGraphClient, subscriptions: azureResource.AzureResourceSubscription[], resourceQuery: string): Promise { const allResources: T[] = []; let totalProcessed = 0; let doQuery = async (skipToken?: string) => { const response = await resourceClient.resources({ - subscriptions: [subId], + subscriptions: subscriptions.map(subscription => subscription.id), query: resourceQuery, options: { resultFormat: 'objectArray', @@ -118,10 +118,10 @@ export abstract class ResourceServiceBase { + public async getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: msRest.ServiceClientCredentials, account: azdata.Account): Promise { const convertedResources: U[] = []; const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint }); - let graphResources = await queryGraphResources(resourceClient, subscription.id, this.query); + let graphResources = await queryGraphResources(resourceClient, subscriptions, this.query); let ids = new Set(); graphResources.forEach((res) => { if (!ids.has(res.id)) { diff --git a/extensions/azurecore/src/azureResource/utils.ts b/extensions/azurecore/src/azureResource/utils.ts index 937bfd9498..db5411b153 100644 --- a/extensions/azurecore/src/azureResource/utils.ts +++ b/extensions/azurecore/src/azureResource/utils.ts @@ -121,7 +121,7 @@ export async function getResourceGroups(appContext: AppContext, account?: azdata const token = tokenResponse.token; const tokenType = tokenResponse.tokenType; - result.resourceGroups.push(...await service.getResources(subscription, new TokenCredentials(token, tokenType), account)); + result.resourceGroups.push(...await service.getResources([subscription], new TokenCredentials(token, tokenType), account)); } catch (err) { const error = new Error(localize('azure.accounts.getResourceGroups.queryError', "Error fetching resource groups for account {0} ({1}) subscription {2} ({3}) tenant {4} : {5}", account.displayInfo.displayName, diff --git a/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts b/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts index 05558665ec..d08112c2f9 100644 --- a/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts +++ b/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts @@ -109,7 +109,7 @@ describe('AzureResourceDatabaseTreeDataProvider.getChildren', function (): void mockExtensionContext = TypeMoq.Mock.ofType(); sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken)); - mockDatabaseService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabases)); + mockDatabaseService.setup((o) => o.getResources([mockSubscription], TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabases)); mockExtensionContext.setup((o) => o.asAbsolutePath(TypeMoq.It.isAnyString())).returns(() => TypeMoq.It.isAnyString()); }); diff --git a/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts b/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts index 9ede7f79ad..310f6145ae 100644 --- a/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts +++ b/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts @@ -108,7 +108,7 @@ describe('AzureResourceDatabaseServerTreeDataProvider.getChildren', function (): mockExtensionContext = TypeMoq.Mock.ofType(); sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken)); - mockDatabaseServerService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabaseServers)); + mockDatabaseServerService.setup((o) => o.getResources([mockSubscription], TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabaseServers)); mockExtensionContext.setup((o) => o.asAbsolutePath(TypeMoq.It.isAnyString())).returns(() => TypeMoq.It.isAnyString()); });