walmart block query restricted account

src/Common/Constants.ts

@@ -177,6 +177,7 @@ export class HttpHeaders {
   public static activityId: string = "x-ms-activity-id";
   public static apiType: string = "x-ms-cosmos-apitype";
   public static authorization: string = "authorization";
+  public static graphAuthorization: string = "graph-authorization";
   public static collectionIndexTransformationProgress: string =
     "x-ms-documentdb-collection-index-transformation-progress";
   public static continuation: string = "x-ms-continuation";

src/Common/dataAccess/readCollections.ts

@@ -35,21 +35,14 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
       }
     }
 
-    try {
-      const responses = await Promise.all(promises);
-      responses.forEach((response) => {
-        collections.push(response.resource as DataModels.Collection);
-      });
+    const responses = await Promise.all(promises);
+    responses.forEach((response) => {
+      collections.push(response.resource as DataModels.Collection);
+    });
 
-      // Sort collections by id before returning
-      collections.sort((a, b) => a.id.localeCompare(b.id));
-      return collections;
-    } catch (error) {
-      handleError(error, "ReadCollections", `Error while querying containers for database ${databaseId}`);
-      throw error;
-    } finally {
-      clearMessage();
-    }
+    // Sort collections by id before returning
+    collections.sort((a, b) => a.id.localeCompare(b.id));
+    return collections;
   }
 
   try {

src/Common/dataAccess/readDatabases.ts

@@ -22,13 +22,6 @@ export async function readDatabases(): Promise<DataModels.Database[]> {
     for (const collectionResourceId in tokensData.resourceTokens) {
       // Dictionary key looks like this: dbs/SampleDB/colls/Container
       const resourceIdObj = collectionResourceId.split("/");
-
-      if (resourceIdObj.length !== 4) {
-        handleError(`Resource key not recognized: ${resourceIdObj}`, "ReadDatabases", `Error while querying databases`);
-        clearMessage();
-        return [];
-      }
-
       const databaseId = resourceIdObj[1];
 
       databaseIdsSet.add(databaseId);
@@ -44,7 +37,6 @@ export async function readDatabases(): Promise<DataModels.Database[]> {
         id: databaseId,
         collections: [],
       }));
-    clearMessage();
     return databases;
   }
 

src/Contracts/AzureResourceGraph.ts

@@ -1,13 +0,0 @@
-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/Explorer/Quickstart/VCoreMongoQuickstartCommands.ts

@@ -1,6 +1,5 @@
 export const newDbAndCollectionCommand = `use quickstartDB
-db.createCollection('sampleCollection')
-`;
+db.createCollection('sampleCollection')`;
 
 export const newDbAndCollectionCommandForDisplay = `use quickstartDB // Create new database named 'quickstartDB' or switch to it if it already exists
 
@@ -17,25 +16,19 @@ export const loadDataCommand = `db.sampleCollection.insertMany([
     {title: "War and Peace", author: "Leo Tolstoy", pages: 1392},
     {title: "The Odyssey", author: "Homer", pages: 374},
     {title: "Ulysses", author: "James Joyce", pages: 730}
-  ])
-`;
+  ])`;
 
-export const findOrwellCommand = `db.sampleCollection.find({author: "George Orwell"})
-`;
+export const queriesCommand = `db.sampleCollection.find({author: "George Orwell"})
 
-export const findOrwellCommandForDisplay = `// Query to find all books written by "George Orwell"
-db.sampleCollection.find({author: "George Orwell"})`;
-
-export const findByPagesCommand = `db.sampleCollection.find({pages: {$gt: 500}})
-`;
-
-export const findByPagesCommandForDisplay = `// Query to find all books with more than 500 pages
 db.sampleCollection.find({pages: {$gt: 500}})
-`;
 
-export const findAndSortCommand = `db.sampleCollection.find({}).sort({pages: 1})
-`;
+db.sampleCollection.find({}).sort({pages: 1})`;
 
-export const findAndSortCommandForDisplay = `// Query to find all books and sort them by the number of pages in ascending order
-db.sampleCollection.find({}).sort({pages: 1})
-`;
+export const queriesCommandForDisplay = `// Query to find all books written by "George Orwell"
+db.sampleCollection.find({author: "George Orwell"})
+
+// Query to find all books with more than 500 pages
+db.sampleCollection.find({pages: {$gt: 500}})
+
+// Query to find all books and sort them by the number of pages in ascending order
+db.sampleCollection.find({}).sort({pages: 1})`;

src/Explorer/Quickstart/VCoreMongoQuickstartGuide.tsx

@@ -11,15 +11,11 @@ import {
 } from "@fluentui/react";
 import { customPivotHeaderRenderer } from "Explorer/Quickstart/Shared/QuickstartRenderUtilities";
 import {
-  findAndSortCommand,
-  findAndSortCommandForDisplay,
-  findByPagesCommand,
-  findByPagesCommandForDisplay,
-  findOrwellCommand,
-  findOrwellCommandForDisplay,
   loadDataCommand,
   newDbAndCollectionCommand,
   newDbAndCollectionCommandForDisplay,
+  queriesCommand,
+  queriesCommandForDisplay,
 } from "Explorer/Quickstart/VCoreMongoQuickstartCommands";
 import { useTerminal } from "hooks/useTerminal";
 import React, { useState } from "react";
@@ -194,17 +190,17 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                 </Text>
                 <DefaultButton
                   style={{ marginTop: 16, width: 110 }}
-                  onClick={() => useTerminal.getState().sendMessage(findOrwellCommand)}
+                  onClick={() => useTerminal.getState().sendMessage(queriesCommand)}
                 >
                   Try query
                 </DefaultButton>
                 <Stack horizontal style={{ marginTop: 16 }}>
                   <TextField
-                    id="findOrwellCommand"
+                    id="queriesCommand"
                     multiline
-                    rows={2}
+                    rows={5}
                     readOnly
-                    defaultValue={findOrwellCommandForDisplay}
+                    defaultValue={queriesCommandForDisplay}
                     styles={{
                       root: { width: "90%" },
                       field: {
@@ -218,65 +214,7 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                     iconProps={{
                       iconName: "Copy",
                     }}
-                    onClick={() => onCopyBtnClicked("#findOrwellCommand")}
-                  />
-                </Stack>
-                <DefaultButton
-                  style={{ marginTop: 32, width: 110 }}
-                  onClick={() => useTerminal.getState().sendMessage(findByPagesCommand)}
-                >
-                  Try query
-                </DefaultButton>
-                <Stack horizontal style={{ marginTop: 16 }}>
-                  <TextField
-                    id="findByPagesCommand"
-                    multiline
-                    rows={2}
-                    readOnly
-                    defaultValue={findByPagesCommandForDisplay}
-                    styles={{
-                      root: { width: "90%" },
-                      field: {
-                        backgroundColor: "#EEEEEE",
-                        fontFamily:
-                          "Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New",
-                      },
-                    }}
-                  />
-                  <IconButton
-                    iconProps={{
-                      iconName: "Copy",
-                    }}
-                    onClick={() => onCopyBtnClicked("#findByPagesCommand")}
-                  />
-                </Stack>
-                <DefaultButton
-                  style={{ marginTop: 32, width: 110 }}
-                  onClick={() => useTerminal.getState().sendMessage(findAndSortCommand)}
-                >
-                  Try query
-                </DefaultButton>
-                <Stack horizontal style={{ marginTop: 16 }}>
-                  <TextField
-                    id="findAndSortCommand"
-                    multiline
-                    rows={2}
-                    readOnly
-                    defaultValue={findAndSortCommandForDisplay}
-                    styles={{
-                      root: { width: "90%" },
-                      field: {
-                        backgroundColor: "#EEEEEE",
-                        fontFamily:
-                          "Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New",
-                      },
-                    }}
-                  />
-                  <IconButton
-                    iconProps={{
-                      iconName: "Copy",
-                    }}
-                    onClick={() => onCopyBtnClicked("#findAndSortCommand")}
+                    onClick={() => onCopyBtnClicked("#queriesCommand")}
                   />
                 </Stack>
               </Stack>
@@ -298,7 +236,7 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                   hosted in the cloud, to Azure Cosmos DB for MongoDB vCore.&nbsp;
                   <Link
                     target="_blank"
-                    href="https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/migration-options"
+                    href="https://learn.microsoft.com/azure-data-studio/extensions/azure-cosmos-db-mongodb-extension"
                   >
                     Learn more
                   </Link>

src/HostedExplorer.tsx

@@ -51,6 +51,7 @@ const App: React.FunctionComponent = () => {
           authType: AuthType.AAD,
           databaseAccount,
           authorizationToken: armToken,
+          graphAuthorizationToken: graphToken
         };
       } else if (authType === AuthType.EncryptedToken) {
         frameWindow.hostedConfig = {

src/HostedExplorerChildFrame.ts

@@ -10,6 +10,7 @@ export interface AAD {
   authType: AuthType.AAD;
   databaseAccount: DatabaseAccount;
   authorizationToken: string;
+  graphAuthorizationToken: string;
 }
 
 export interface ConnectionString {

src/Platform/Hosted/extractFeatures.ts

@@ -14,7 +14,6 @@ export type Features = {
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
   readonly enableAadDataPlane: boolean;
-  readonly enableResourceGraph: boolean;
   readonly enableKoResourceTree: boolean;
   readonly hostedDataExplorer: boolean;
   readonly junoEndpoint?: string;
@@ -74,7 +73,6 @@ 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/Terminal/JupyterLabAppFactory.ts

@@ -10,13 +10,9 @@ import { userContext } from "UserContext";
 export class JupyterLabAppFactory {
   private isShellStarted: boolean | undefined;
   private checkShellStarted: ((content: string | undefined) => void) | undefined;
-  private onShellExited: (restartShell: boolean) => void;
-  private restartShell: boolean;
+  private onShellExited: () => void;
 
   private isShellExited(content: string | undefined) {
-    if (userContext.apiType === "VCoreMongo" && content?.includes("MongoServerError: Invalid key")) {
-      this.restartShell = true;
-    }
     return content?.includes("cosmosuser@");
   }
 
@@ -36,11 +32,10 @@ export class JupyterLabAppFactory {
     this.isShellStarted = content?.includes("Enter password");
   }
 
-  constructor(closeTab: (restartShell: boolean) => void) {
+  constructor(closeTab: () => void) {
     this.onShellExited = closeTab;
     this.isShellStarted = false;
     this.checkShellStarted = undefined;
-    this.restartShell = false;
 
     switch (userContext.apiType) {
       case "Mongo":
@@ -74,7 +69,7 @@ export class JupyterLabAppFactory {
         if (!this.isShellStarted) {
           this.checkShellStarted(content);
         } else if (this.isShellExited(content)) {
-          this.onShellExited(this.restartShell);
+          this.onShellExited();
         }
       }
     }, this);

src/Terminal/index.ts

@@ -11,8 +11,6 @@ import { JupyterLabAppFactory } from "./JupyterLabAppFactory";
 import { TerminalProps } from "./TerminalProps";
 import "./index.css";
 
-let session: ITerminalConnection | undefined;
-
 const createServerSettings = (props: TerminalProps): ServerConnection.ISettings => {
   let body: BodyInit | undefined;
   let headers: HeadersInit | undefined;
@@ -51,7 +49,7 @@ const createServerSettings = (props: TerminalProps): ServerConnection.ISettings
   return ServerConnection.makeSettings(options);
 };
 
-const initTerminal = async (props: TerminalProps): Promise<void> => {
+const initTerminal = async (props: TerminalProps): Promise<ITerminalConnection | undefined> => {
   // Initialize userContext (only properties which are needed by TelemetryProcessor)
   updateUserContext({
     subscriptionId: props.subscriptionId,
@@ -61,37 +59,28 @@ const initTerminal = async (props: TerminalProps): Promise<void> => {
   });
 
   const serverSettings = createServerSettings(props);
-
-  createTerminalApp(props, serverSettings);
-};
-
-const createTerminalApp = async (props: TerminalProps, serverSettings: ServerConnection.ISettings) => {
   const data = { baseUrl: serverSettings.baseUrl };
   const startTime = TelemetryProcessor.traceStart(Action.OpenTerminal, data);
 
   try {
-    session = await new JupyterLabAppFactory((restartShell: boolean) =>
-      closeTab(props, serverSettings, restartShell),
-    ).createTerminalApp(serverSettings);
+    const session = await new JupyterLabAppFactory(() => closeTab(props.tabId)).createTerminalApp(serverSettings);
     TelemetryProcessor.traceSuccess(Action.OpenTerminal, data, startTime);
+    return session;
   } catch (error) {
     TelemetryProcessor.traceFailure(Action.OpenTerminal, data, startTime);
-    session = undefined;
+    return undefined;
   }
 };
 
-const closeTab = (props: TerminalProps, serverSettings: ServerConnection.ISettings, restartShell: boolean): void => {
-  if (restartShell) {
-    createTerminalApp(props, serverSettings);
-  } else {
-    window.parent.postMessage(
-      { type: MessageTypes.CloseTab, data: { tabId: props.tabId }, signature: "pcIframe" },
-      window.document.referrer,
-    );
-  }
+const closeTab = (tabId: string): void => {
+  window.parent.postMessage(
+    { type: MessageTypes.CloseTab, data: { tabId: tabId }, signature: "pcIframe" },
+    window.document.referrer,
+  );
 };
 
 const main = async (): Promise<void> => {
+  let session: ITerminalConnection | undefined;
   postRobot.on(
     "props",
     {
@@ -102,7 +91,7 @@ const main = async (): Promise<void> => {
       // Typescript definition for event is wrong. So read props by casting to <any>
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const props = (event as any).data as TerminalProps;
-      await initTerminal(props);
+      session = await initTerminal(props);
     },
   );
 

src/UserContext.ts

@@ -79,6 +79,7 @@ interface UserContext {
   collectionCreationDefaults: CollectionCreationDefaults;
   sampleDataConnectionInfo?: ParsedResourceTokenConnectionString;
   readonly vcoreMongoConnectionParams?: VCoreMongoConnectionParams;
+  readonly accountRestrictedFromUser?: boolean;
 }
 
 export type ApiType = "SQL" | "Mongo" | "Gremlin" | "Tables" | "Cassandra" | "Postgres" | "VCoreMongo";
@@ -171,3 +172,4 @@ function apiType(account: DatabaseAccount | undefined): ApiType {
 }
 
 export { updateUserContext, userContext };
+

src/Utils/AuthorizationUtils.ts

@@ -60,3 +60,27 @@ export function getMsalInstance() {
   const msalInstance = new msal.PublicClientApplication(config);
   return msalInstance;
 }
+
+export async function isAccountRestrictedFromUser(accountName: string, graphToken: string): Promise<boolean> {
+  const checkUserAccessUrl: string = "https://localhost:12901/api/guest/accountrestrictions/accountrestrictedfromuser";
+  // const authorizationHeader = getAuthorizationHeader();
+  try {
+    const response: Response = await fetch(checkUserAccessUrl, {
+      method: "POST",
+      body: JSON.stringify({
+        accountName
+      }),
+      headers: {
+        // [authorizationHeader.header]: authorizationHeader.token,
+        [Constants.HttpHeaders.graphAuthorization]: graphToken,
+        [Constants.HttpHeaders.contentType]: "application/json",
+      }
+    });
+
+    const responseText: string = await response.text();
+    return responseText.toLowerCase() === "true";
+  } catch (e) {
+    console.log(e);
+    throw new Error(e);
+  }
+}

src/hooks/useDatabaseAccounts.tsx

@@ -1,9 +1,6 @@
-import { HttpHeaders } from "Common/Constants";
-import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { DatabaseAccount } from "../Contracts/DataModels";
-/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface AccountListResult {
   nextLink: string;
@@ -33,56 +30,10 @@ 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) => fetchDatabaseAccountsFromGraph(subscriptionId, armToken),
+    (_, subscriptionId, armToken) => fetchDatabaseAccounts(subscriptionId, armToken),
   );
   return data;
 }

src/hooks/useKnockoutExplorer.ts

@@ -36,7 +36,7 @@ import { extractFeatures } from "../Platform/Hosted/extractFeatures";
 import { CollectionCreation } from "../Shared/Constants";
 import { DefaultExperienceUtility } from "../Shared/DefaultExperienceUtility";
 import { Node, PortalEnv, updateUserContext, userContext } from "../UserContext";
-import { getAuthorizationHeader, getMsalInstance } from "../Utils/AuthorizationUtils";
+import { getAuthorizationHeader, getMsalInstance, isAccountRestrictedFromUser } from "../Utils/AuthorizationUtils";
 import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 import { listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
 import { DatabaseAccountListKeysResult } from "../Utils/arm/generatedClients/cosmos/types";
@@ -227,9 +227,11 @@ async function configureHosted(): Promise<Explorer> {
 
 async function configureHostedWithAAD(config: AAD): Promise<Explorer> {
   // TODO: Refactor. updateUserContext needs to be called twice because listKeys below depends on userContext.authorizationToken
+  const accountRestrictedFromUser: boolean = await isAccountRestrictedFromUser(config.databaseAccount.name, config.graphAuthorizationToken);
   updateUserContext({
     authType: AuthType.AAD,
     authorizationToken: `Bearer ${config.authorizationToken}`,
+    accountRestrictedFromUser
   });
   const account = config.databaseAccount;
   const accountResourceId = account.id;

src/hooks/useSubscriptions.tsx

@@ -1,9 +1,6 @@
-import { HttpHeaders } from "Common/Constants";
-import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { Subscription } from "../Contracts/DataModels";
-/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface SubscriptionListResult {
   nextLink: string;
@@ -35,58 +32,10 @@ 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) => fetchSubscriptionsFromGraph(armToken),
+    (_, armToken) => fetchSubscriptions(armToken),
   );
   return data;
 }