mirror of
https://github.com/ckaczor/azuredatastudio.git
synced 2026-01-26 09:35:38 -05:00
Update azure graph queries to allow multiple subs (#12124)
This commit is contained in:
@@ -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[]>;
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)) {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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());
|
||||
});
|
||||
|
||||
|
||||
@@ -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());
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user