Expose the graph API (#11919)

* Exposing the graph API

* Azure managed instance retrival

* Fix compile error

extensions/azurecore/src/azureResource/azure-resource.d.ts

@@ -34,6 +34,16 @@ declare module 'azureResource' {
 			loginName: string;
 		}
 
+		export interface AzureGraphResource extends Omit<AzureResource, 'tenant'> {
+			tenantId: string;
+			type: string;
+			location: string;
+		}
+
+		export interface AzureSqlManagedInstanceResource extends AzureGraphResource {
+
+		}
+
 		export interface AzureResourceResourceGroup extends AzureResource {
 		}
 

extensions/azurecore/src/azureResource/utils.ts

@@ -6,13 +6,14 @@
 import * as azdata from 'azdata';
 import * as nls from 'vscode-nls';
 import { azureResource } from 'azureResource';
-import { GetResourceGroupsResult, GetSubscriptionsResult } from 'azurecore';
+import { GetResourceGroupsResult, GetSubscriptionsResult, ResourceQueryResult } from 'azurecore';
 import { isArray } from 'util';
 import { AzureResourceGroupService } from './providers/resourceGroup/resourceGroupService';
 import { TokenCredentials } from '@azure/ms-rest-js';
 import { AppContext } from '../appContext';
 import { IAzureResourceSubscriptionService } from './interfaces';
 import { AzureResourceServiceNames } from './constants';
+import { ResourceGraphClient } from '@azure/arm-resourcegraph';
 
 const localize = nls.loadMessageBundle();
 
@@ -139,6 +140,64 @@ 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) {
+	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"));
+		if (!ignoreErrors) {
+			throw error;
+		}
+		result.errors.push(error);
+		return result;
+	}
+
+	if (!subscription.tenant) {
+		const error = new Error(localize('azure.accounts.runResourceQuery.errors.noTenantSpecifiedForSubscription', "Invalid tenant for subscription"));
+		if (!ignoreErrors) {
+			throw error;
+		}
+		result.errors.push(error);
+		return result;
+	}
+
+	const tokenResponse = await azdata.accounts.getAccountSecurityToken(account, subscription.tenant, azdata.AzureResource.ResourceManagement);
+	const token = tokenResponse.token;
+	const tokenType = tokenResponse.tokenType;
+	const credential = new TokenCredentials(token, tokenType);
+
+	const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
+
+	const allResources: T[] = [];
+	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;
+
+}
+
 export async function getSubscriptions(appContext: AppContext, account?: azdata.Account, ignoreErrors: boolean = false): Promise<GetSubscriptionsResult> {
 	const result: GetSubscriptionsResult = { subscriptions: [], errors: [] };
 	if (!account?.properties?.tenants || !isArray(account.properties.tenants)) {

extensions/azurecore/src/azurecore.d.ts

@@ -71,8 +71,12 @@ 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>>;
 	}
 
 	export type GetSubscriptionsResult = { subscriptions: azureResource.AzureResourceSubscription[], errors: Error[] };
 	export type GetResourceGroupsResult = { resourceGroups: azureResource.AzureResourceResourceGroup[], errors: Error[] };
+
+	export type ResourceQueryResult<T extends azureResource.AzureGraphResource> = { resources: T[], errors: Error[] };
 }

extensions/azurecore/src/extension.ts

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