Update azure graph queries to allow multiple subs (#12124)

This commit is contained in:
Charles Gagnon
2020-09-03 19:07:37 -07:00
committed by GitHub
parent b1ce53ade3
commit 9004769865
6 changed files with 13 additions and 13 deletions

View File

@@ -38,5 +38,5 @@ export interface IAzureResourceNodeWithProviderId {
}
export interface IAzureResourceService<T extends azureResource.AzureResource> {
getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials, account: Account): Promise<T[]>;
getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: msRest.ServiceClientCredentials, account: Account): Promise<T[]>;
}

View File

@@ -15,13 +15,13 @@ interface DatabaseGraphData extends GraphData {
kind: string;
}
export class AzureResourceDatabaseService implements IAzureResourceService<azureResource.AzureResourceDatabase> {
public async getResources(subscription: azureResource.AzureResourceSubscription, credential: ServiceClientCredentials, account: Account): Promise<azureResource.AzureResourceDatabase[]> {
public async getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: ServiceClientCredentials, account: Account): Promise<azureResource.AzureResourceDatabase[]> {
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<GraphData>(resourceClient, subscription.id, serversQuery);
let dbQueryPromise = queryGraphResources<GraphData>(resourceClient, subscription.id, 'where type == "microsoft.sql/servers/databases"');
let serverQueryPromise = queryGraphResources<GraphData>(resourceClient, subscriptions, serversQuery);
let dbQueryPromise = queryGraphResources<GraphData>(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<azure
});
// Match database ID. When calling exec [0] is full match, [1] is resource group name, [2] is server name
const svrIdRegExp = new RegExp(`\/subscriptions\/${subscription.id}\/resourceGroups\/(.+)\/providers\/Microsoft\.Sql\/servers\/(.+)\/databases\/.+`);
const svrIdRegExp = new RegExp(`\/subscriptions\/.+\/resourceGroups\/(.+)\/providers\/Microsoft\.Sql\/servers\/(.+)\/databases\/.+`);
dbByGraph.forEach(db => {
// Filter master DBs, and for all others find their server to get login info

View File

@@ -44,7 +44,7 @@ export abstract class ResourceTreeDataProviderBase<T extends azureResource.Azure
const response = await azdata.accounts.getAccountSecurityToken(element.account, element.tenantId, azdata.AzureResource.ResourceManagement);
const credential = new msRest.TokenCredentials(response.token, response.tokenType);
const resources: T[] = await this._resourceService.getResources(element.subscription, credential, element.account) || <T[]>[];
const resources: T[] = await this._resourceService.getResources([element.subscription], credential, element.account) || <T[]>[];
return resources;
}
@@ -63,12 +63,12 @@ export interface GraphData {
}
export async function queryGraphResources<T extends GraphData>(resourceClient: ResourceGraphClient, subId: string, resourceQuery: string): Promise<T[]> {
export async function queryGraphResources<T extends GraphData>(resourceClient: ResourceGraphClient, subscriptions: azureResource.AzureResourceSubscription[], resourceQuery: string): Promise<T[]> {
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<T extends GraphData, U extends azureRe
*/
protected abstract get query(): string;
public async getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials, account: azdata.Account): Promise<U[]> {
public async getResources(subscriptions: azureResource.AzureResourceSubscription[], credential: msRest.ServiceClientCredentials, account: azdata.Account): Promise<U[]> {
const convertedResources: U[] = [];
const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
let graphResources = await queryGraphResources<T>(resourceClient, subscription.id, this.query);
let graphResources = await queryGraphResources<T>(resourceClient, subscriptions, this.query);
let ids = new Set<string>();
graphResources.forEach((res) => {
if (!ids.has(res.id)) {

View File

@@ -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,

View File

@@ -109,7 +109,7 @@ describe('AzureResourceDatabaseTreeDataProvider.getChildren', function (): void
mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();
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());
});

View File

@@ -108,7 +108,7 @@ describe('AzureResourceDatabaseServerTreeDataProvider.getChildren', function ():
mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();
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());
});