Fetch subs and accounts via MS graph (#1677)

* fetch subs and accounts via graph

* fixed subscription rendering

* add feature flag enableResourceGraph

* add feature flag enableResourceGraph

---------

Co-authored-by: Asier Isayas <aisayas@microsoft.com>

src/Contracts/AzureResourceGraph.ts

@@ -0,0 +1,13 @@
+export interface QueryRequestOptions {
+  $skipToken?: string;
+  $top?: number;
+  subscriptions: string[];
+}
+
+export interface QueryResponse {
+  $skipToken: string;
+  count: number;
+  data: any;
+  resultTruncated: boolean;
+  totalRecords: number;
+}

src/Contracts/DataModels.ts

@@ -88,13 +88,13 @@ export interface GenerateTokenResponse {
 }
 
 export interface Subscription {
-  uniqueDisplayName: string;
+  uniqueDisplayName?: string;
   displayName: string;
   subscriptionId: string;
-  tenantId: string;
+  tenantId?: string;
   state: string;
-  subscriptionPolicies: SubscriptionPolicies;
-  authorizationSource: string;
+  subscriptionPolicies?: SubscriptionPolicies;
+  authorizationSource?: string;
 }
 
 export interface SubscriptionPolicies {

src/HostedExplorer.tsx

@@ -7,9 +7,6 @@ import "../less/hostedexplorer.less";
 import { AuthType } from "./AuthType";
 import { DatabaseAccount } from "./Contracts/DataModels";
 import "./Explorer/Menus/NavBar/MeControlComponent.less";
-import { useAADAuth } from "./hooks/useAADAuth";
-import { useConfig } from "./hooks/useConfig";
-import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 import { HostedExplorerChildFrame } from "./HostedExplorerChildFrame";
 import { AccountSwitcher } from "./Platform/Hosted/Components/AccountSwitcher";
 import { ConnectExplorer } from "./Platform/Hosted/Components/ConnectExplorer";
@@ -20,6 +17,9 @@ import { SignInButton } from "./Platform/Hosted/Components/SignInButton";
 import "./Platform/Hosted/ConnectScreen.less";
 import { extractMasterKeyfromConnectionString } from "./Platform/Hosted/HostedUtils";
 import "./Shared/appInsights";
+import { useAADAuth } from "./hooks/useAADAuth";
+import { useConfig } from "./hooks/useConfig";
+import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 
 initializeIcons();
 

src/Platform/Hosted/extractFeatures.ts

@@ -14,6 +14,7 @@ export type Features = {
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
   readonly enableAadDataPlane: boolean;
+  readonly enableResourceGraph: boolean;
   readonly enableKoResourceTree: boolean;
   readonly hostedDataExplorer: boolean;
   readonly junoEndpoint?: string;
@@ -73,6 +74,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     canExceedMaximumValue: "true" === get("canexceedmaximumvalue"),
     cosmosdb: "true" === get("cosmosdb"),
     enableAadDataPlane: "true" === get("enableaaddataplane"),
+    enableResourceGraph: "true" === get("enableresourcegraph"),
     enableChangeFeedPolicy: "true" === get("enablechangefeedpolicy"),
     enableFixedCollectionWithSharedThroughput: "true" === get("enablefixedcollectionwithsharedthroughput"),
     enableKOPanel: "true" === get("enablekopanel"),

src/hooks/useDatabaseAccounts.tsx

@@ -1,6 +1,10 @@
+import { HttpHeaders } from "Common/Constants";
+import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
+import { userContext } from "UserContext";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { DatabaseAccount } from "../Contracts/DataModels";
+/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface AccountListResult {
   nextLink: string;
@@ -30,10 +34,59 @@ export async function fetchDatabaseAccounts(subscriptionId: string, accessToken:
   return accounts.sort((a, b) => a.name.localeCompare(b.name));
 }
 
+export async function fetchDatabaseAccountsFromGraph(
+  subscriptionId: string,
+  accessToken: string,
+): Promise<DatabaseAccount[]> {
+  const headers = new Headers();
+  const bearer = `Bearer ${accessToken}`;
+
+  headers.append("Authorization", bearer);
+  headers.append(HttpHeaders.contentType, "application/json");
+  const databaseAccountsQuery = "resources | where type =~ 'microsoft.documentdb/databaseaccounts'";
+  const apiVersion = "2021-03-01";
+  const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+
+  const databaseAccounts: DatabaseAccount[] = [];
+  let skipToken: string;
+  do {
+    const body = {
+      query: databaseAccountsQuery,
+      subscriptions: [subscriptionId],
+      ...(skipToken && {
+        options: {
+          $skipToken: skipToken,
+        } as QueryRequestOptions,
+      }),
+    };
+
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(body),
+    });
+
+    if (!response.ok) {
+      throw new Error(await response.text());
+    }
+
+    const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
+    skipToken = queryResponse.$skipToken;
+    queryResponse.data?.map((databaseAccount: any) => {
+      databaseAccounts.push(databaseAccount as DatabaseAccount);
+    });
+  } while (skipToken);
+
+  return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
+}
+
 export function useDatabaseAccounts(subscriptionId: string, armToken: string): DatabaseAccount[] | undefined {
   const { data } = useSWR(
     () => (armToken && subscriptionId ? ["databaseAccounts", subscriptionId, armToken] : undefined),
-    (_, subscriptionId, armToken) => fetchDatabaseAccounts(subscriptionId, armToken),
+    (_, subscriptionId, armToken) =>
+      userContext.features.enableResourceGraph
+        ? fetchDatabaseAccountsFromGraph(subscriptionId, armToken)
+        : fetchDatabaseAccounts(subscriptionId, armToken),
   );
   return data;
 }

src/hooks/useSubscriptions.tsx

@@ -1,6 +1,10 @@
+import { HttpHeaders } from "Common/Constants";
+import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
+import { userContext } from "UserContext";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { Subscription } from "../Contracts/DataModels";
+/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface SubscriptionListResult {
   nextLink: string;
@@ -32,10 +36,59 @@ export async function fetchSubscriptions(accessToken: string): Promise<Subscript
   return subscriptions.sort((a, b) => a.displayName.localeCompare(b.displayName));
 }
 
+export async function fetchSubscriptionsFromGraph(accessToken: string): Promise<Subscription[]> {
+  const headers = new Headers();
+  const bearer = `Bearer ${accessToken}`;
+
+  headers.append("Authorization", bearer);
+  headers.append(HttpHeaders.contentType, "application/json");
+  const subscriptionsQuery =
+    "resources | where type == 'microsoft.documentdb/databaseaccounts' | join kind=inner ( resourcecontainers | where type == 'microsoft.resources/subscriptions' | project subscriptionId, subscriptionName = name, subscriptionState = tostring(parse_json(properties).state) ) on subscriptionId |  summarize by subscriptionId, subscriptionName, subscriptionState";
+  const apiVersion = "2021-03-01";
+  const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+
+  const subscriptions: Subscription[] = [];
+  let skipToken: string;
+  do {
+    const body = {
+      query: subscriptionsQuery,
+      ...(skipToken && {
+        options: {
+          $skipToken: skipToken,
+        } as QueryRequestOptions,
+      }),
+    };
+
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(body),
+    });
+
+    if (!response.ok) {
+      throw new Error(await response.text());
+    }
+
+    const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
+    skipToken = queryResponse.$skipToken;
+
+    queryResponse.data?.map((subscription: any) => {
+      subscriptions.push({
+        displayName: subscription.subscriptionName,
+        subscriptionId: subscription.subscriptionId,
+        state: subscription.subscriptionState,
+      } as Subscription);
+    });
+  } while (skipToken);
+
+  return subscriptions.sort((a, b) => a.displayName.localeCompare(b.displayName));
+}
+
 export function useSubscriptions(armToken: string): Subscription[] | undefined {
   const { data } = useSWR(
     () => (armToken ? ["subscriptions", armToken] : undefined),
-    (_, armToken) => fetchSubscriptions(armToken),
+    (_, armToken) =>
+      userContext.features.enableResourceGraph ? fetchSubscriptionsFromGraph(armToken) : fetchSubscriptions(armToken),
   );
   return data;
 }