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);
+		}
 	};
 }
 

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

@@ -8,6 +8,9 @@ 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>> {
+		throw new Error('Method not implemented.');
+	}
 	getSubscriptions(_account?: azdata.Account | undefined, _ignoreErrors?: boolean | undefined): Thenable<azurecore.GetSubscriptionsResult> {
 		throw new Error('Method not implemented.');
 	}

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

@@ -0,0 +1,35 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as azdata from 'azdata';
+import * as azurecore from 'azurecore';
+import { azureResource } from 'azureResource';
+
+async function getAzureCoreAPI(): Promise<azurecore.IExtension> {
+	const api = (await vscode.extensions.getExtension(azurecore.extension.name)?.activate()) as azurecore.IExtension;
+	if (!api) {
+		throw new Error('azure core API undefined for sql-migration');
+	}
+	return api;
+}
+
+export type Subscription = azureResource.AzureResourceSubscription;
+export async function getSubscriptions(account: azdata.Account): Promise<Subscription[]> {
+	const api = await getAzureCoreAPI();
+	const subscriptions = await api.getSubscriptions(account, false);
+
+	return subscriptions.subscriptions;
+}
+
+export type AzureProduct = azureResource.AzureGraphResource;
+export type SqlManagedInstance = azureResource.AzureSqlManagedInstanceResource;
+export async function getAvailableManagedInstanceProducts(account: azdata.Account, subscription: Subscription): Promise<SqlManagedInstance[]> {
+	const api = await getAzureCoreAPI();
+
+	const result = await api.runGraphQuery<azureResource.AzureSqlManagedInstanceResource>(account, subscription, false, 'where type == "microsoft.sql/managedinstances"');
+
+	return result.resources;
+}

extensions/sql-migration/src/typings/ref.d.ts

@@ -7,4 +7,5 @@
 /// <reference path='../../../../src/sql/azdata.d.ts'/>
 /// <reference path='../../../../src/sql/azdata.proposed.d.ts'/>
 /// <reference path='../../../azurecore/src/azurecore.d.ts'/>
+/// <reference path='../../../azurecore/src/azureResource/azure-resource.d.ts'/>
 /// <reference types='@types/node'/>

extensions/sql-migration/src/wizard/subscriptionSelectionPage.ts

@@ -8,11 +8,16 @@ import { MigrationWizardPage } from '../models/migrationWizardPage';
 import { MigrationStateModel, StateChangeEvent } from '../models/stateMachine';
 import { SUBSCRIPTION_SELECTION_PAGE_TITLE, SUBSCRIPTION_SELECTION_AZURE_ACCOUNT_TITLE, SUBSCRIPTION_SELECTION_AZURE_PRODUCT_TITLE, SUBSCRIPTION_SELECTION_AZURE_SUBSCRIPTION_TITLE } from '../models/strings';
 import { Disposable } from 'vscode';
+import { getSubscriptions, Subscription, getAvailableManagedInstanceProducts, AzureProduct } from '../api/azure';
 
-interface AccountValue extends azdata.CategoryValue {
+interface GenericValue<T> extends azdata.CategoryValue {
-	account: azdata.Account;
+	value: T;
 }
 
+type AccountValue = GenericValue<azdata.Account>;
+type SubscriptionValue = GenericValue<Subscription>;
+type ProductValue = GenericValue<AzureProduct>;
+
 export class SubscriptionSelectionPage extends MigrationWizardPage {
 	private disposables: Disposable[] = [];
 
@@ -50,6 +55,10 @@ export class SubscriptionSelectionPage extends MigrationWizardPage {
 			values: [],
 		});
 
+		this.disposables.push(dropDown.component().onValueChanged(() => {
+			this.accountValueChanged().catch(console.error);
+		}));
+
 		return {
 			component: dropDown.component(),
 			title: SUBSCRIPTION_SELECTION_AZURE_ACCOUNT_TITLE
@@ -60,7 +69,10 @@ export class SubscriptionSelectionPage extends MigrationWizardPage {
 		const dropDown = view.modelBuilder.dropDown().withProperties<azdata.DropDownProperties>({
 			values: [],
 		});
-		this.setupSubscriptionListener();
+
+		this.disposables.push(dropDown.component().onValueChanged(() => {
+			this.subscriptionValueChanged().catch(console.error);
+		}));
 
 		return {
 			component: dropDown.component(),
@@ -72,7 +84,6 @@ export class SubscriptionSelectionPage extends MigrationWizardPage {
 		const dropDown = view.modelBuilder.dropDown().withProperties<azdata.DropDownProperties>({
 			values: [],
 		});
-		this.setupProductListener();
 
 		return {
 			component: dropDown.component(),
@@ -80,22 +91,34 @@ export class SubscriptionSelectionPage extends MigrationWizardPage {
 		};
 	}
 
-
+	private async accountValueChanged(): Promise<void> {
-	private setupSubscriptionListener(): void {
+		const account = this.getPickedAccount();
-		this.disposables.push(this.accountDropDown!.component.onValueChanged((event) => {
+		if (account) {
-			console.log(event);
+			const subscriptions = await getSubscriptions(account);
-		}));
+			await this.populateSubscriptionValues(subscriptions);
+		}
 	}
 
-	private setupProductListener(): void {
+	private async subscriptionValueChanged(): Promise<void> {
-		this.disposables.push(this.subscriptionDropDown!.component.onValueChanged((event) => {
+		const account = this.getPickedAccount();
-			console.log(event);
+		const subscription = this.getPickedSubscription();
-		}));
+
+		const results = await getAvailableManagedInstanceProducts(account!, subscription!);
+
+		this.populateProductValues(results);
+	}
+
+	private getPickedAccount(): azdata.Account | undefined {
+		const accountValue: AccountValue | undefined = this.accountDropDown?.component.value as AccountValue;
+		return accountValue?.value;
+	}
+
+	private getPickedSubscription(): Subscription | undefined {
+		const accountValue: SubscriptionValue | undefined = this.subscriptionDropDown?.component.value as SubscriptionValue;
+		return accountValue?.value;
 	}
 
 	private async populateAccountValues(): Promise<void> {
-
-
 		let accounts = await azdata.accounts.getAllAccounts();
 		accounts = accounts.filter(a => a.key.providerId.startsWith('azure') && !a.isStale);
 
@@ -103,11 +126,37 @@ export class SubscriptionSelectionPage extends MigrationWizardPage {
 			return {
 				displayName: a.displayInfo.displayName,
 				name: a.key.accountId,
-				account: a
+				value: a
 			};
 		});
 
 		this.accountDropDown!.component.values = values;
+		await this.accountValueChanged();
+	}
+
+	private async populateSubscriptionValues(subscriptions: Subscription[]): Promise<void> {
+		const values: SubscriptionValue[] = subscriptions.map(sub => {
+			return {
+				displayName: sub.name,
+				name: sub.id,
+				value: sub
+			};
+		});
+
+		this.subscriptionDropDown!.component.values = values;
+		await this.subscriptionValueChanged();
+	}
+
+	private async populateProductValues(products: AzureProduct[]) {
+		const values: ProductValue[] = products.map(prod => {
+			return {
+				displayName: prod.name,
+				name: prod.id,
+				value: prod
+			};
+		});
+
+		this.productDropDown!.component.values = values;
 	}
 
 	public async onPageEnter(): Promise<void> {

extensions/sql-migration/src/wizard/wizardController.ts

@@ -53,7 +53,7 @@ export class WizardController {
 			const canEnter = await pages[lastPage]?.canEnter() ?? true;
 
 			return canEnter && canLeave;
-			// return true;
+			// return true
 		});
 
 		await Promise.all(wizardSetupPromises);