Add Sample Data to Resource Tree (#1499)

* Add Sample Data to Resource Tree

* Format

* Fix strict build

* Fix lint

* Fixed implementation to show Sample data container

* Udated logic based on TokenCollection

* Re-configure copilot flag

---------

Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>

src/Common/SampleDataClient.ts

@@ -0,0 +1,26 @@
+import * as Cosmos from "@azure/cosmos";
+import { userContext } from "UserContext";
+
+let _sampleDataclient: Cosmos.CosmosClient;
+
+export function sampleDataClient(): Cosmos.CosmosClient {
+  if (_sampleDataclient) {
+    return _sampleDataclient;
+  }
+
+  const sampleDataConnectionInfo = userContext.sampleDataConnectionInfo;
+  const options: Cosmos.CosmosClientOptions = {
+    endpoint: sampleDataConnectionInfo.accountEndpoint,
+    tokenProvider: async () => {
+      const sampleDataConnectionInfo = userContext.sampleDataConnectionInfo;
+      return Promise.resolve(sampleDataConnectionInfo.resourceToken);
+    },
+    connectionPolicy: {
+      enableEndpointDiscovery: false,
+    },
+    userAgentSuffix: "Azure Portal",
+  };
+
+  _sampleDataclient = new Cosmos.CosmosClient(options);
+  return _sampleDataclient;
+}

src/Common/dataAccess/readCollection.ts

@@ -1,13 +1,39 @@
+import { CosmosClient } from "@azure/cosmos";
+import { sampleDataClient } from "Common/SampleDataClient";
+import { userContext } from "UserContext";
 import * as DataModels from "../../Contracts/DataModels";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { client } from "../CosmosClient";
 import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 
 export async function readCollection(databaseId: string, collectionId: string): Promise<DataModels.Collection> {
+  const cosmosClient = client();
+  return await readCollectionInternal(cosmosClient, databaseId, collectionId);
+}
+
+export async function readSampleCollection(): Promise<DataModels.Collection> {
+  const cosmosClient = sampleDataClient();
+
+  const sampleDataConnectionInfo = userContext.sampleDataConnectionInfo;
+  const databaseId = sampleDataConnectionInfo?.databaseId;
+  const collectionId = sampleDataConnectionInfo?.collectionId;
+
+  if (!databaseId || !collectionId) {
+    return undefined;
+  }
+
+  return await readCollectionInternal(cosmosClient, databaseId, collectionId);
+}
+
+export async function readCollectionInternal(
+  cosmosClient: CosmosClient,
+  databaseId: string,
+  collectionId: string
+): Promise<DataModels.Collection> {
   let collection: DataModels.Collection;
   const clearMessage = logConsoleProgress(`Querying container ${collectionId}`);
   try {
-    const response = await client().database(databaseId).container(collectionId).read();
+    const response = await cosmosClient.database(databaseId).container(collectionId).read();
     collection = response.resource as DataModels.Collection;
   } catch (error) {
     handleError(error, "ReadCollection", `Error while querying container ${collectionId}`);

src/Explorer/Explorer.tsx

@@ -1,39 +1,39 @@
 import { Link } from "@fluentui/react/lib/Link";
 import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import { IGalleryItem } from "Juno/JunoClient";
+import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
-import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
 import shallow from "zustand/shallow";
 import { AuthType } from "../AuthType";
 import { BindingHandlersRegisterer } from "../Bindings/BindingHandlersRegisterer";
 import * as Constants from "../Common/Constants";
 import { Areas, ConnectionStatusType, HttpStatusCodes, Notebook, PoolIdType } from "../Common/Constants";
-import { readCollection } from "../Common/dataAccess/readCollection";
-import { readDatabases } from "../Common/dataAccess/readDatabases";
 import { getErrorMessage, getErrorStack, handleError } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
 import { QueriesClient } from "../Common/QueriesClient";
+import { readCollection, readSampleCollection } from "../Common/dataAccess/readCollection";
+import { readDatabases } from "../Common/dataAccess/readDatabases";
 import * as DataModels from "../Contracts/DataModels";
 import { ContainerConnectionInfo, IPhoenixServiceInfo, IProvisionData, IResponse } from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 import { GitHubOAuthService } from "../GitHub/GitHubOAuthService";
-import { useSidePanel } from "../hooks/useSidePanel";
-import { useTabs } from "../hooks/useTabs";
 import { PhoenixClient } from "../Phoenix/PhoenixClient";
 import * as ExplorerSettings from "../Shared/ExplorerSettings";
 import { Action, ActionModifiers } from "../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../UserContext";
 import { getCollectionName, getUploadName } from "../Utils/APITypeUtils";
-import { update } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { listByDatabaseAccount } from "../Utils/arm/generatedClients/cosmosNotebooks/notebookWorkspaces";
 import { stringToBlob } from "../Utils/BlobUtils";
 import { isCapabilityEnabled } from "../Utils/CapabilityUtils";
 import { fromContentUri, toRawContentUri } from "../Utils/GitHubUtils";
 import * as NotificationConsoleUtils from "../Utils/NotificationConsoleUtils";
 import { logConsoleError, logConsoleInfo, logConsoleProgress } from "../Utils/NotificationConsoleUtils";
+import { update } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { listByDatabaseAccount } from "../Utils/arm/generatedClients/cosmosNotebooks/notebookWorkspaces";
+import { useSidePanel } from "../hooks/useSidePanel";
+import { useTabs } from "../hooks/useTabs";
 import "./ComponentRegisterer";
 import { DialogProps, useDialog } from "./Controls/Dialog";
 import { GalleryTab as GalleryTabKind } from "./Controls/NotebookGallery/GalleryViewerComponent";
@@ -1272,5 +1272,26 @@ export default class Explorer {
     if (useNotebook.getState().isPhoenixNotebooks) {
       await this.initNotebooks(userContext.databaseAccount);
     }
+
+    await this.refreshSampleData();
+  }
+
+  public async refreshSampleData(): Promise<void> {
+    if (!userContext.sampleDataConnectionInfo) {
+      return;
+    }
+
+    const collection: DataModels.Collection = await readSampleCollection();
+    if (!collection) {
+      return;
+    }
+
+    const databaseId = userContext.sampleDataConnectionInfo?.databaseId;
+    if (!databaseId) {
+      return;
+    }
+
+    const sampleDataResourceTokenCollection = new ResourceTokenCollection(this, databaseId, collection);
+    useDatabases.setState({ sampleDataResourceTokenCollection });
   }
 }

src/Explorer/Tree/ResourceTree.tsx

@@ -1,14 +1,15 @@
 import { Callout, DirectionalHint, ICalloutProps, ILinkProps, Link, Stack, Text } from "@fluentui/react";
-import * as React from "react";
+import { SampleDataTree } from "Explorer/Tree/SampleDataTree";
 import { getItemName } from "Utils/APITypeUtils";
+import * as React from "react";
 import shallow from "zustand/shallow";
 import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
-import DeleteIcon from "../../../images/delete.svg";
 import GalleryIcon from "../../../images/GalleryIcon.svg";
-import FileIcon from "../../../images/notebook/file-cosmos.svg";
+import DeleteIcon from "../../../images/delete.svg";
 import CopyIcon from "../../../images/notebook/Notebook-copy.svg";
 import NewNotebookIcon from "../../../images/notebook/Notebook-new.svg";
 import NotebookIcon from "../../../images/notebook/Notebook-resource.svg";
+import FileIcon from "../../../images/notebook/file-cosmos.svg";
 import PublishIcon from "../../../images/notebook/publish_content.svg";
 import RefreshIcon from "../../../images/refresh-cosmos.svg";
 import CollectionIcon from "../../../images/tree-collection.svg";
@@ -16,14 +17,14 @@ import { Areas, ConnectionStatusType, Notebook } from "../../Common/Constants";
 import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
-import { useSidePanel } from "../../hooks/useSidePanel";
-import { useTabs } from "../../hooks/useTabs";
 import { LocalStorageUtility, StorageKey } from "../../Shared/StorageUtility";
 import { Action, ActionModifiers, Source } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
 import { isServerlessAccount } from "../../Utils/CapabilityUtils";
 import * as GitHubUtils from "../../Utils/GitHubUtils";
+import { useSidePanel } from "../../hooks/useSidePanel";
+import { useTabs } from "../../hooks/useTabs";
 import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
 import { AccordionComponent, AccordionItemComponent } from "../Controls/Accordion/AccordionComponent";
 import { useDialog } from "../Controls/Dialog";
@@ -764,9 +765,15 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
   };
 
   const dataRootNode = buildDataTree();
+  const isSampleDataEnabled = userContext.sampleDataConnectionInfo && userContext.apiType === "SQL";
+  const sampleDataResourceTokenCollection = useDatabases((state) => state.sampleDataResourceTokenCollection);
 
-  if (isNotebookEnabled) {
   return (
+    <>
+      {!isNotebookEnabled && !isSampleDataEnabled && (
+        <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
+      )}
+      {isNotebookEnabled && !isSampleDataEnabled && (
         <>
           <AccordionComponent>
             <AccordionItemComponent title={"DATA"} isExpanded={!gitHubNotebooksContentRoot}>
@@ -779,8 +786,38 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
 
           {buildGalleryCallout()}
         </>
-    );
+      )}
-  }
+      {!isNotebookEnabled && isSampleDataEnabled && (
+        <>
+          <AccordionComponent>
+            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
+              <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
+            </AccordionItemComponent>
+            <AccordionItemComponent title={"SAMPLE DATA"}>
+              <SampleDataTree sampleDataResourceTokenCollection={sampleDataResourceTokenCollection} />
+            </AccordionItemComponent>
+          </AccordionComponent>
 
-  return <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />;
+          {buildGalleryCallout()}
+        </>
+      )}
+      {isNotebookEnabled && isSampleDataEnabled && (
+        <>
+          <AccordionComponent>
+            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
+              <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
+            </AccordionItemComponent>
+            <AccordionItemComponent title={"SAMPLE DATA"}>
+              <SampleDataTree sampleDataResourceTokenCollection={sampleDataResourceTokenCollection} />
+            </AccordionItemComponent>
+            <AccordionItemComponent title={"NOTEBOOKS"}>
+              <TreeComponent className="notebookResourceTree" rootNode={buildNotebooksTree()} />
+            </AccordionItemComponent>
+          </AccordionComponent>
+
+          {buildGalleryCallout()}
+        </>
+      )}
+    </>
+  );
 };

src/Explorer/Tree/SampleDataTree.tsx

@@ -0,0 +1,41 @@
+import React, { useEffect, useState } from "react";
+import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
+import CollectionIcon from "../../../images/tree-collection.svg";
+import * as ViewModels from "../../Contracts/ViewModels";
+import { TreeComponent, TreeNode } from "../Controls/TreeComponent/TreeComponent";
+
+export const SampleDataTree = ({
+  sampleDataResourceTokenCollection,
+}: {
+  sampleDataResourceTokenCollection: ViewModels.CollectionBase;
+}): JSX.Element => {
+  const [root, setRoot] = useState<TreeNode | undefined>(undefined);
+
+  useEffect(() => {
+    if (sampleDataResourceTokenCollection) {
+      const updatedSampleTree: TreeNode = {
+        label: sampleDataResourceTokenCollection.databaseId,
+        isExpanded: true,
+        iconSrc: CosmosDBIcon,
+        children: [
+          {
+            label: sampleDataResourceTokenCollection.id(),
+            iconSrc: CollectionIcon,
+            isExpanded: true,
+            className: "collectionHeader",
+            children: [
+              {
+                label: "Items",
+              },
+            ],
+          },
+        ],
+      };
+      setRoot(updatedSampleTree);
+    }
+  }, [sampleDataResourceTokenCollection]);
+
+  return (
+    <TreeComponent className="sampleDataResourceTree" rootNode={root || { label: "Sample data not initialized." }} />
+  );
+};

src/Explorer/useDatabases.ts

@@ -8,6 +8,7 @@ import { useSelectedNode } from "./useSelectedNode";
 interface DatabasesState {
   databases: ViewModels.Database[];
   resourceTokenCollection: ViewModels.CollectionBase;
+  sampleDataResourceTokenCollection: ViewModels.CollectionBase;
   updateDatabase: (database: ViewModels.Database) => void;
   addDatabases: (databases: ViewModels.Database[]) => void;
   deleteDatabase: (database: ViewModels.Database) => void;
@@ -28,6 +29,7 @@ interface DatabasesState {
 export const useDatabases: UseStore<DatabasesState> = create((set, get) => ({
   databases: [],
   resourceTokenCollection: undefined,
+  sampleDataResourceTokenCollection: undefined,
   updateDatabase: (updatedDatabase: ViewModels.Database) =>
     set((state) => {
       const updatedDatabases = state.databases.map((database: ViewModels.Database) => {

src/Platform/Hosted/Helpers/ResourceTokenUtils.ts

@@ -1,17 +1,18 @@
 export interface ParsedResourceTokenConnectionString {
-  accountEndpoint: string;
+  accountEndpoint?: string;
-  collectionId: string;
+  collectionId?: string;
-  databaseId: string;
+  databaseId?: string;
   partitionKey?: string;
-  resourceToken: string;
+  resourceToken?: string;
 }
 
 export function parseResourceTokenConnectionString(connectionString: string): ParsedResourceTokenConnectionString {
-  let accountEndpoint: string;
+  let accountEndpoint: string | undefined;
-  let collectionId: string;
+  let collectionId: string | undefined;
-  let databaseId: string;
+  let databaseId: string | undefined;
-  let partitionKey: string;
+  let partitionKey: string | undefined;
-  let resourceToken: string;
+  let resourceToken: string | undefined;
+
   const connectionStringParts = connectionString.split(";");
   connectionStringParts.forEach((part: string) => {
     if (part.startsWith("type=resource")) {
@@ -39,5 +40,5 @@ export function parseResourceTokenConnectionString(connectionString: string): Pa
 }
 
 export function isResourceTokenConnectionString(connectionString: string): boolean {
-  return connectionString && connectionString.includes("type=resource");
+  return !!connectionString && connectionString.includes("type=resource");
 }

src/UserContext.ts

@@ -1,5 +1,6 @@
 import { useCarousel } from "hooks/useCarousel";
 import { usePostgres } from "hooks/usePostgres";
+import { ParsedResourceTokenConnectionString } from "Platform/Hosted/Helpers/ResourceTokenUtils";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
 import { AuthType } from "./AuthType";
@@ -69,6 +70,7 @@ interface UserContext {
   readonly postgresConnectionStrParams?: PostgresConnectionStrParams;
   readonly isReplica?: boolean;
   collectionCreationDefaults: CollectionCreationDefaults;
+  sampleDataConnectionInfo?: ParsedResourceTokenConnectionString;
 }
 
 export type ApiType = "SQL" | "Mongo" | "Gremlin" | "Tables" | "Cassandra" | "Postgres";
@@ -157,4 +159,4 @@ function apiType(account: DatabaseAccount | undefined): ApiType {
   return "SQL";
 }
 
-export { userContext, updateUserContext };
+export { updateUserContext, userContext };

src/hooks/useKnockoutExplorer.ts

@@ -1,13 +1,13 @@
+import { createUri } from "Common/UrlUtility";
 import Explorer from "Explorer/Explorer";
+import { getNetworkSettingsWarningMessage } from "Utils/NetworkUtility";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
 import { useEffect, useState } from "react";
-import { getNetworkSettingsWarningMessage } from "Utils/NetworkUtility";
-import { applyExplorerBindings } from "../applyExplorerBindings";
 import { AuthType } from "../AuthType";
 import { AccountKind, Flights } from "../Common/Constants";
 import { normalizeArmEndpoint } from "../Common/EnvironmentUtility";
 import { sendMessage, sendReadyMessage } from "../Common/MessageHandler";
-import { configContext, Platform, updateConfigContext } from "../ConfigContext";
+import { Platform, configContext, updateConfigContext } from "../ConfigContext";
 import { ActionType, DataExplorerAction } from "../Contracts/ActionContracts";
 import { MessageTypes } from "../Contracts/ExplorerContracts";
 import { DataExplorerInputsFrame } from "../Contracts/ViewModels";
@@ -21,19 +21,20 @@ import {
   ResourceToken,
 } from "../HostedExplorerChildFrame";
 import { emulatorAccount } from "../Platform/Emulator/emulatorAccount";
-import { extractFeatures } from "../Platform/Hosted/extractFeatures";
 import { parseResourceTokenConnectionString } from "../Platform/Hosted/Helpers/ResourceTokenUtils";
 import {
   getDatabaseAccountKindFromExperience,
   getDatabaseAccountPropertiesFromMetadata,
 } from "../Platform/Hosted/HostedUtils";
+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 { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 import { listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
 import { DatabaseAccountListKeysResult } from "../Utils/arm/generatedClients/cosmos/types";
-import { getMsalInstance } from "../Utils/AuthorizationUtils";
+import { applyExplorerBindings } from "../applyExplorerBindings";
-import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 
 // This hook will create a new instance of Explorer.ts and bind it to the DOM
 // This hook has a LOT of magic, but ideally we can delete it once we have removed KO and switched entirely to React
@@ -61,6 +62,10 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
           setExplorer(explorer);
         }
       }
+
+      if (userContext.features.enableCopilot) {
+        await updateContextForSampleData();
+      }
     };
     effect();
   }, [platform]);
@@ -409,3 +414,29 @@ interface PortalMessage {
   type?: MessageTypes;
   inputs?: DataExplorerInputsFrame;
 }
+
+async function updateContextForSampleData(): Promise<void> {
+  if (!userContext.features.enableCopilot) {
+    return;
+  }
+
+  const url = createUri(`${configContext.BACKEND_ENDPOINT}`, `/api/tokens/sampledataconnection`);
+  const authorizationHeader = getAuthorizationHeader();
+  const headers = { [authorizationHeader.header]: authorizationHeader.token };
+
+  const response = await window.fetch(url, {
+    headers,
+  });
+
+  if (!response.ok) {
+    return undefined;
+  }
+
+  const data: SampledataconnectionResponse = await response.json();
+  const sampleDataConnectionInfo = parseResourceTokenConnectionString(data.connectionString);
+  updateUserContext({ sampleDataConnectionInfo });
+}
+
+interface SampledataconnectionResponse {
+  connectionString: string;
+}