Create new ResourceTree based on FluentUI Tree (#1603)

* Alternate tree running fluentui v9 Tree component

* Fix tree update after sp, udf and trigger load

* Enable scrolling for subtrees

* Clean up duplicates

* Restore current tree

* Reformat

* Update package-lock.json

package.json

@@ -13,6 +13,7 @@
     "@babel/plugin-proposal-class-properties": "7.12.1",
     "@babel/plugin-proposal-decorators": "7.12.12",
     "@fluentui/react": "8.14.3",
+    "@fluentui/react-components": "9.30.1",
     "@jupyterlab/services": "6.0.2",
     "@jupyterlab/terminal": "3.0.3",
     "@microsoft/applicationinsights-web": "2.6.1",

src/Common/ResourceTreeContainer.tsx

@@ -1,10 +1,10 @@
+import { ResourceTree } from "Explorer/Tree/ResourceTree";
 import React, { FunctionComponent, MutableRefObject, useEffect, useRef } from "react";
 import arrowLeftImg from "../../images/imgarrowlefticon.svg";
 import refreshImg from "../../images/refresh-cosmos.svg";
 import { AuthType } from "../AuthType";
 import Explorer from "../Explorer/Explorer";
 import { ResourceTokenTree } from "../Explorer/Tree/ResourceTokenTree";
-import { ResourceTree } from "../Explorer/Tree/ResourceTree";
 import { userContext } from "../UserContext";
 import { getApiShortDisplayName } from "../Utils/APITypeUtils";
 import { NormalizedEventKey } from "./Constants";
@@ -78,6 +78,8 @@ export const ResourceTreeContainer: FunctionComponent<ResourceTreeContainerProps
           <div style={{ overflowY: "auto" }} data-bind="react:resourceTree" />
         ) : (
           <ResourceTree container={container} />
+          // Uncomment the following line to use the fluent ui tree
+          // <ResourceTree2 container={container} />
         )}
       </div>
       {/*  Collections Window - End */}

src/Explorer/ContextMenuButtonFactory.tsx

@@ -103,7 +103,6 @@ export const createCollectionContextMenuButton = (
     items.push({
       iconSrc: AddStoredProcedureIcon,
       onClick: () => {
-        const selectedCollection: ViewModels.Collection = useSelectedNode.getState().findSelectedCollection();
         selectedCollection && selectedCollection.onNewStoredProcedureClick(selectedCollection, undefined);
       },
       label: "New Stored Procedure",
@@ -112,7 +111,6 @@ export const createCollectionContextMenuButton = (
     items.push({
       iconSrc: AddUdfIcon,
       onClick: () => {
-        const selectedCollection: ViewModels.Collection = useSelectedNode.getState().findSelectedCollection();
         selectedCollection && selectedCollection.onNewUserDefinedFunctionClick(selectedCollection);
       },
       label: "New UDF",
@@ -121,7 +119,6 @@ export const createCollectionContextMenuButton = (
     items.push({
       iconSrc: AddTriggerIcon,
       onClick: () => {
-        const selectedCollection: ViewModels.Collection = useSelectedNode.getState().findSelectedCollection();
         selectedCollection && selectedCollection.onNewTriggerClick(selectedCollection, undefined);
       },
       label: "New Trigger",
@@ -130,13 +127,15 @@ export const createCollectionContextMenuButton = (
 
   items.push({
     iconSrc: DeleteCollectionIcon,
-    onClick: () =>
+    onClick: () => {
+      useSelectedNode.getState().setSelectedNode(selectedCollection);
       useSidePanel
         .getState()
         .openSidePanel(
           "Delete " + getCollectionName(),
           <DeleteCollectionConfirmationPane refreshDatabases={() => container.refreshAllDatabases()} />
-        ),
+        );
+    },
     label: `Delete ${getCollectionName()}`,
     styleClass: "deleteCollectionMenuItem",
   });

src/Explorer/Controls/TreeComponent2/TreeNode2Component.tsx

@@ -0,0 +1,142 @@
+import {
+  Button,
+  Menu,
+  MenuItem,
+  MenuList,
+  MenuPopover,
+  MenuTrigger,
+  Spinner,
+  Tree,
+  TreeItem,
+  TreeItemLayout,
+} from "@fluentui/react-components";
+import { MoreHorizontal20Regular } from "@fluentui/react-icons";
+import * as React from "react";
+
+export interface TreeNode2MenuItem {
+  label: string;
+  onClick: () => void;
+  iconSrc?: string;
+  isDisabled?: boolean;
+  styleClass?: string;
+}
+
+export interface TreeNode2 {
+  label: string;
+  id?: string;
+  children?: TreeNode2[];
+  contextMenu?: TreeNode2MenuItem[];
+  iconSrc?: string;
+  // isExpanded?: boolean;
+  className?: string;
+  isAlphaSorted?: boolean;
+  // data?: any; // Piece of data corresponding to this node
+  timestamp?: number;
+  isLeavesParentsSeparate?: boolean; // Display parents together first, then leaves
+  isLoading?: boolean;
+  isScrollable?: boolean;
+  isSelected?: () => boolean;
+  onClick?: () => void; // Only if a leaf, other click will expand/collapse
+  onExpanded?: () => void;
+  onCollapsed?: () => void;
+  onContextMenuOpen?: () => void;
+}
+
+export interface TreeNode2ComponentProps {
+  node: TreeNode2;
+  className?: string;
+  treeNodeId: string;
+  globalOpenIds: string[];
+}
+
+const getTreeIcon = (iconSrc: string): JSX.Element => <img src={iconSrc} alt="" style={{ width: 20, height: 20 }} />;
+
+export const TreeNode2Component: React.FC<TreeNode2ComponentProps> = ({
+  node,
+  treeNodeId,
+  globalOpenIds,
+}: TreeNode2ComponentProps): JSX.Element => {
+  // const defaultOpenItems = node.isExpanded ? children?.map((child: TreeNode2) => child.label) : undefined;
+  const [isExpanded, setIsExpanded] = React.useState<boolean>(false);
+
+  // Compute whether node is expanded
+  React.useEffect(() => {
+    const isNowExpanded = globalOpenIds && globalOpenIds.includes(treeNodeId);
+    if (!isExpanded && isNowExpanded) {
+      // Catch the transition non-expanded to expanded
+      node.onExpanded?.();
+    }
+    setIsExpanded(isNowExpanded);
+  }, [globalOpenIds, treeNodeId, node, isExpanded]);
+
+  const getSortedChildren = (treeNode: TreeNode2): TreeNode2[] => {
+    if (!treeNode || !treeNode.children) {
+      return undefined;
+    }
+
+    const compareFct = (a: TreeNode2, b: TreeNode2) => a.label.localeCompare(b.label);
+
+    let unsortedChildren;
+    if (treeNode.isLeavesParentsSeparate) {
+      // Separate parents and leave
+      const parents: TreeNode2[] = treeNode.children.filter((node) => node.children);
+      const leaves: TreeNode2[] = treeNode.children.filter((node) => !node.children);
+
+      if (treeNode.isAlphaSorted) {
+        parents.sort(compareFct);
+        leaves.sort(compareFct);
+      }
+
+      unsortedChildren = parents.concat(leaves);
+    } else {
+      unsortedChildren = treeNode.isAlphaSorted ? treeNode.children.sort(compareFct) : treeNode.children;
+    }
+
+    return unsortedChildren;
+  };
+
+  return (
+    <TreeItem value={treeNodeId} itemType={node.children !== undefined ? "branch" : "leaf"} style={{ height: "100%" }}>
+      <TreeItemLayout
+        className={node.className}
+        actions={
+          node.contextMenu && (
+            <Menu>
+              <MenuTrigger disableButtonEnhancement>
+                <Button aria-label="More options" appearance="subtle" icon={<MoreHorizontal20Regular />} />
+              </MenuTrigger>
+              <MenuPopover>
+                <MenuList>
+                  {node.contextMenu.map((menuItem) => (
+                    <MenuItem disabled={menuItem.isDisabled} key={menuItem.label} onClick={menuItem.onClick}>
+                      {menuItem.label}
+                    </MenuItem>
+                  ))}
+                </MenuList>
+              </MenuPopover>
+            </Menu>
+          )
+        }
+        expandIcon={node.isLoading ? <Spinner size="extra-tiny" /> : undefined}
+        iconBefore={node.iconSrc && getTreeIcon(node.iconSrc)}
+      >
+        <span onClick={() => node.onClick?.()}>{node.label}</span>
+      </TreeItemLayout>
+      {!node.isLoading && node.children?.length > 0 && (
+        <Tree
+          // defaultOpenItems={defaultOpenItems}
+          style={{ overflow: node.isScrollable ? "auto" : undefined }}
+        >
+          {getSortedChildren(node).map((childNode: TreeNode2) => (
+            <TreeNode2Component
+              key={childNode.label}
+              node={childNode}
+              treeNodeId={`${treeNodeId}/${childNode.label}`}
+              globalOpenIds={globalOpenIds}
+            />
+          ))}
+        </Tree>
+      )}
+    </TreeItem>
+  );
+};

src/Explorer/Panes/CopyNotebookPane/CopyNotebookPaneComponent.tsx

@@ -9,11 +9,11 @@ import {
   Stack,
   Text,
 } from "@fluentui/react";
+import { GitHubReposTitle } from "Explorer/Tree/ResourceTree";
 import React, { FormEvent, FunctionComponent } from "react";
 import { IPinnedRepo } from "../../../Juno/JunoClient";
 import * as GitHubUtils from "../../../Utils/GitHubUtils";
 import { useNotebook } from "../../Notebook/useNotebook";
-import { ResourceTreeAdapter } from "../../Tree/ResourceTreeAdapter";
 
 interface Location {
   type: "MyNotebooks" | "GitHub";
@@ -65,7 +65,7 @@ export const CopyNotebookPaneComponent: FunctionComponent<CopyNotebookPaneProps>
 
       options.push({
         key: "GitHub-Header",
-        text: ResourceTreeAdapter.GitHubReposTitle,
+        text: GitHubReposTitle,
         itemType: SelectableOptionMenuItemType.Header,
       });
 

src/Explorer/Tree2/ResourceTree.tsx

@@ -0,0 +1,108 @@
+import {
+  BrandVariants,
+  FluentProvider,
+  Theme,
+  Tree,
+  TreeItemValue,
+  TreeOpenChangeData,
+  TreeOpenChangeEvent,
+  createLightTheme,
+} from "@fluentui/react-components";
+import { TreeNode2Component } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
+import { useDatabaseTreeNodes } from "Explorer/Tree2/useDatabaseTreeNodes";
+import * as React from "react";
+import shallow from "zustand/shallow";
+import Explorer from "../Explorer";
+import { useNotebook } from "../Notebook/useNotebook";
+
+export const MyNotebooksTitle = "My Notebooks";
+export const GitHubReposTitle = "GitHub repos";
+
+interface ResourceTreeProps {
+  container: Explorer;
+}
+
+const cosmosdb: BrandVariants = {
+  10: "#020305",
+  20: "#111723",
+  30: "#16263D",
+  40: "#193253",
+  50: "#1B3F6A",
+  60: "#1B4C82",
+  70: "#18599B",
+  80: "#1267B4",
+  90: "#3174C2",
+  100: "#4F82C8",
+  110: "#6790CF",
+  120: "#7D9ED5",
+  130: "#92ACDC",
+  140: "#A6BAE2",
+  150: "#BAC9E9",
+  160: "#CDD8EF",
+};
+
+const lightTheme: Theme = {
+  ...createLightTheme(cosmosdb),
+};
+
+export const DATA_TREE_LABEL = "DATA";
+
+/**
+ * Top-level tree that has no label, but contains all subtrees
+ */
+export const ResourceTree2: React.FC<ResourceTreeProps> = ({ container }: ResourceTreeProps): JSX.Element => {
+  const {
+    isNotebookEnabled,
+    // myNotebooksContentRoot,
+    // galleryContentRoot,
+    // gitHubNotebooksContentRoot,
+    // updateNotebookItem,
+  } = useNotebook(
+    (state) => ({
+      isNotebookEnabled: state.isNotebookEnabled,
+      myNotebooksContentRoot: state.myNotebooksContentRoot,
+      galleryContentRoot: state.galleryContentRoot,
+      gitHubNotebooksContentRoot: state.gitHubNotebooksContentRoot,
+      updateNotebookItem: state.updateNotebookItem,
+    }),
+    shallow
+  );
+  // const { activeTab } = useTabs();
+  const databaseTreeNodes = useDatabaseTreeNodes(container, isNotebookEnabled);
+  const dataNodeTree = {
+    id: "data",
+    label: DATA_TREE_LABEL,
+    isExpanded: true,
+    className: "accordionItemHeader",
+    children: databaseTreeNodes,
+    isScrollable: true,
+  };
+
+  const [openItems, setOpenItems] = React.useState<Iterable<TreeItemValue>>([DATA_TREE_LABEL]);
+
+  const handleOpenChange = (event: TreeOpenChangeEvent, data: TreeOpenChangeData) => setOpenItems(data.openItems);
+
+  return (
+    <>
+      <FluentProvider theme={lightTheme} style={{ overflow: "hidden" }}>
+        <Tree
+          aria-label="CosmosDB resources"
+          openItems={openItems}
+          onOpenChange={handleOpenChange}
+          size="small"
+          style={{ height: "100%" }}
+        >
+          {[dataNodeTree].map((node) => (
+            <TreeNode2Component
+              key={node.label}
+              className="dataResourceTree"
+              node={node}
+              treeNodeId={node.label}
+              globalOpenIds={[...openItems].map((item) => item.toString())}
+            />
+          ))}
+        </Tree>
+      </FluentProvider>
+    </>
+  );
+};

src/Explorer/Tree2/containerTreeNodeUtil.ts

@@ -0,0 +1,297 @@
+import { TreeNode2 } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
+import TabsBase from "Explorer/Tabs/TabsBase";
+import StoredProcedure from "Explorer/Tree/StoredProcedure";
+import Trigger from "Explorer/Tree/Trigger";
+import UserDefinedFunction from "Explorer/Tree/UserDefinedFunction";
+import { useDatabases } from "Explorer/useDatabases";
+import { getItemName } from "Utils/APITypeUtils";
+import { isServerlessAccount } from "Utils/CapabilityUtils";
+import CollectionIcon from "../../../images/tree-collection.svg";
+import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
+import * as DataModels from "../../Contracts/DataModels";
+import * as ViewModels from "../../Contracts/ViewModels";
+import { userContext } from "../../UserContext";
+import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
+import Explorer from "../Explorer";
+import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
+import { mostRecentActivity } from "../MostRecentActivity/MostRecentActivity";
+import { useNotebook } from "../Notebook/useNotebook";
+import { useSelectedNode } from "../useSelectedNode";
+
+export const buildCollectionNode = (
+  database: ViewModels.Database,
+  collection: ViewModels.Collection,
+  isNotebookEnabled: boolean,
+  container: Explorer,
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void
+): TreeNode2 => {
+  const showScriptNodes = userContext.apiType === "SQL" || userContext.apiType === "Gremlin";
+  const children: TreeNode2[] = [];
+  children.push({
+    label: getItemName(),
+    id: collection.isSampleCollection ? "sampleItems" : "",
+    onClick: () => {
+      collection.openTab();
+      // push to most recent
+      mostRecentActivity.collectionWasOpened(userContext.databaseAccount?.id, collection);
+    },
+    isSelected: () =>
+      useSelectedNode
+        .getState()
+        .isDataNodeSelected(collection.databaseId, collection.id(), [
+          ViewModels.CollectionTabKind.Documents,
+          ViewModels.CollectionTabKind.Graph,
+        ]),
+    contextMenu: ResourceTreeContextMenuButtonFactory.createCollectionContextMenuButton(container, collection),
+  });
+
+  if (
+    isNotebookEnabled &&
+    userContext.apiType === "Mongo" &&
+    isPublicInternetAccessAllowed() &&
+    useNotebook.getState().isPhoenixFeatures
+  ) {
+    children.push({
+      label: "Schema (Preview)",
+      onClick: collection.onSchemaAnalyzerClick.bind(collection),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.SchemaAnalyzer]),
+    });
+  }
+
+  if (userContext.apiType !== "Cassandra" || !isServerlessAccount()) {
+    let id = "";
+    if (collection.isSampleCollection) {
+      id = database.isDatabaseShared() ? "sampleSettings" : "sampleScaleSettings";
+    }
+
+    children.push({
+      id,
+      label: database.isDatabaseShared() || isServerlessAccount() ? "Settings" : "Scale & Settings",
+      onClick: collection.onSettingsClick.bind(collection),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [
+            ViewModels.CollectionTabKind.CollectionSettingsV2,
+          ]),
+    });
+  }
+
+  const schemaNode: TreeNode2 = buildSchemaNode(collection, container, refreshActiveTab);
+  if (schemaNode) {
+    children.push(schemaNode);
+  }
+
+  const onUpdateDatabase = () => useDatabases.getState().updateDatabase(database);
+
+  if (showScriptNodes) {
+    children.push(buildStoredProcedureNode(collection, container, refreshActiveTab, onUpdateDatabase));
+    children.push(buildUserDefinedFunctionsNode(collection, container, refreshActiveTab, onUpdateDatabase));
+    children.push(buildTriggerNode(collection, container, refreshActiveTab, onUpdateDatabase));
+  }
+
+  // This is a rewrite of showConflicts
+  const showConflicts =
+    userContext?.databaseAccount?.properties.enableMultipleWriteLocations &&
+    collection.rawDataModel &&
+    !!collection.rawDataModel.conflictResolutionPolicy;
+
+  if (showConflicts) {
+    children.push({
+      label: "Conflicts",
+      onClick: collection.onConflictsClick.bind(collection),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Conflicts]),
+    });
+  }
+
+  return {
+    label: collection.id(),
+    iconSrc: CollectionIcon,
+    children: children,
+    className: "collectionHeader",
+    contextMenu: ResourceTreeContextMenuButtonFactory.createCollectionContextMenuButton(container, collection),
+    onClick: () => {
+      useSelectedNode.getState().setSelectedNode(collection);
+    },
+    onExpanded: () => {
+      // Rewritten version of expandCollapseCollection
+      useSelectedNode.getState().setSelectedNode(collection);
+      useCommandBar.getState().setContextButtons([]);
+      refreshActiveTab(
+        (tab: TabsBase) =>
+          tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId
+      );
+    },
+    isSelected: () => useSelectedNode.getState().isDataNodeSelected(collection.databaseId, collection.id()),
+    onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(collection),
+  };
+};
+
+const buildStoredProcedureNode = (
+  collection: ViewModels.Collection,
+  container: Explorer,
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
+  onUpdateDatabase: () => void
+): TreeNode2 => {
+  return {
+    label: "Stored Procedures",
+    children: collection.storedProcedures().map((sp: StoredProcedure) => ({
+      label: sp.id(),
+      onClick: sp.open.bind(sp),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.StoredProcedures]),
+      contextMenu: ResourceTreeContextMenuButtonFactory.createStoreProcedureContextMenuItems(container, sp),
+    })),
+    onExpanded: async () => {
+      await collection.loadStoredProcedures();
+      collection.selectedSubnodeKind(ViewModels.CollectionTabKind.StoredProcedures);
+      refreshActiveTab(
+        (tab: TabsBase) =>
+          tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId
+      );
+      onUpdateDatabase();
+    },
+  };
+};
+
+const buildUserDefinedFunctionsNode = (
+  collection: ViewModels.Collection,
+  container: Explorer,
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
+  onUpdateDatabase: () => void
+): TreeNode2 => {
+  return {
+    label: "User Defined Functions",
+    children: collection.userDefinedFunctions().map((udf: UserDefinedFunction) => ({
+      label: udf.id(),
+      onClick: udf.open.bind(udf),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [
+            ViewModels.CollectionTabKind.UserDefinedFunctions,
+          ]),
+      contextMenu: ResourceTreeContextMenuButtonFactory.createUserDefinedFunctionContextMenuItems(container, udf),
+    })),
+    onExpanded: async () => {
+      await collection.loadUserDefinedFunctions();
+      collection.selectedSubnodeKind(ViewModels.CollectionTabKind.UserDefinedFunctions);
+      refreshActiveTab(
+        (tab: TabsBase) =>
+          tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId
+      );
+      onUpdateDatabase();
+    },
+  };
+};
+
+const buildTriggerNode = (
+  collection: ViewModels.Collection,
+  container: Explorer,
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
+  onUpdateDatabase: () => void
+): TreeNode2 => {
+  return {
+    label: "Triggers",
+    children: collection.triggers().map((trigger: Trigger) => ({
+      label: trigger.id(),
+      onClick: trigger.open.bind(trigger),
+      isSelected: () =>
+        useSelectedNode
+          .getState()
+          .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Triggers]),
+      contextMenu: ResourceTreeContextMenuButtonFactory.createTriggerContextMenuItems(container, trigger),
+    })),
+    onExpanded: async () => {
+      await collection.loadTriggers();
+      collection.selectedSubnodeKind(ViewModels.CollectionTabKind.Triggers);
+      refreshActiveTab(
+        (tab: TabsBase) =>
+          tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId
+      );
+      onUpdateDatabase();
+    },
+  };
+};
+
+const buildSchemaNode = (
+  collection: ViewModels.Collection,
+  container: Explorer,
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void
+): TreeNode2 => {
+  if (collection.analyticalStorageTtl() === undefined) {
+    return undefined;
+  }
+
+  if (!collection.schema || !collection.schema.fields) {
+    return undefined;
+  }
+
+  return {
+    label: "Schema",
+    children: getSchemaNodes(collection.schema.fields),
+    onClick: () => {
+      collection.selectedSubnodeKind(ViewModels.CollectionTabKind.Schema);
+      refreshActiveTab((tab: TabsBase) => tab.collection && tab.collection.rid === collection.rid);
+    },
+  };
+};
+
+const getSchemaNodes = (fields: DataModels.IDataField[]): TreeNode2[] => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const schema: any = {};
+
+  //unflatten
+  fields.forEach((field: DataModels.IDataField) => {
+    const path: string[] = field.path.split(".");
+    const fieldProperties = [field.dataType.name, `HasNulls: ${field.hasNulls}`];
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    let current: any = {};
+    path.forEach((name: string, pathIndex: number) => {
+      if (pathIndex === 0) {
+        if (schema[name] === undefined) {
+          if (pathIndex === path.length - 1) {
+            schema[name] = fieldProperties;
+          } else {
+            schema[name] = {};
+          }
+        }
+        current = schema[name];
+      } else {
+        if (current[name] === undefined) {
+          if (pathIndex === path.length - 1) {
+            current[name] = fieldProperties;
+          } else {
+            current[name] = {};
+          }
+        }
+        current = current[name];
+      }
+    });
+  });
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const traverse = (obj: any): TreeNode2[] => {
+    const children: TreeNode2[] = [];
+
+    if (obj !== undefined && !Array.isArray(obj) && typeof obj === "object") {
+      Object.entries(obj).forEach(([key, value]) => {
+        children.push({ label: key, children: traverse(value) });
+      });
+    } else if (Array.isArray(obj)) {
+      return [{ label: obj[0] }, { label: obj[1] }];
+    }
+
+    return children;
+  };
+
+  return traverse(schema);
+};

src/Explorer/Tree2/useDatabaseTreeNodes.ts

@@ -0,0 +1,83 @@
+import { TreeNode2 } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
+import TabsBase from "Explorer/Tabs/TabsBase";
+import { buildCollectionNode } from "Explorer/Tree2/containerTreeNodeUtil";
+import { useDatabases } from "Explorer/useDatabases";
+import { useTabs } from "hooks/useTabs";
+import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
+import * as ViewModels from "../../Contracts/ViewModels";
+import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
+import Explorer from "../Explorer";
+import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
+import { useSelectedNode } from "../useSelectedNode";
+
+export const useDatabaseTreeNodes = (container: Explorer, isNotebookEnabled: boolean): TreeNode2[] => {
+  const databases = useDatabases((state) => state.databases);
+  const { refreshActiveTab } = useTabs();
+
+  const databaseTreeNodes: TreeNode2[] = databases.map((database: ViewModels.Database) => {
+    const databaseNode: TreeNode2 = {
+      label: database.id(),
+      iconSrc: CosmosDBIcon,
+      className: "databaseHeader",
+      children: [],
+      isSelected: () => useSelectedNode.getState().isDataNodeSelected(database.id()),
+      contextMenu: ResourceTreeContextMenuButtonFactory.createDatabaseContextMenu(container, database.id()),
+      onExpanded: async () => {
+        useSelectedNode.getState().setSelectedNode(database);
+        if (databaseNode.children?.length === 0) {
+          databaseNode.isLoading = true;
+        }
+        await database.expandDatabase();
+        databaseNode.isLoading = false;
+        useCommandBar.getState().setContextButtons([]);
+        refreshActiveTab((tab: TabsBase) => tab.collection?.databaseId === database.id());
+      },
+      onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(database),
+    };
+
+    if (database.isDatabaseShared()) {
+      databaseNode.children.push({
+        id: database.isSampleDB ? "sampleScaleSettings" : "",
+        label: "Scale",
+        isSelected: () =>
+          useSelectedNode
+            .getState()
+            .isDataNodeSelected(database.id(), undefined, [ViewModels.CollectionTabKind.DatabaseSettingsV2]),
+        onClick: database.onSettingsClick.bind(database),
+      });
+    }
+
+    // Find collections
+    database
+      .collections()
+      .forEach((collection: ViewModels.Collection) =>
+        databaseNode.children.push(
+          buildCollectionNode(database, collection, isNotebookEnabled, container, refreshActiveTab)
+        )
+      );
+
+    if (database.collectionsContinuationToken) {
+      const loadMoreNode: TreeNode2 = {
+        label: "load more",
+        className: "loadMoreHeader",
+        onClick: async () => {
+          await database.loadCollections();
+          useDatabases.getState().updateDatabase(database);
+        },
+      };
+      databaseNode.children.push(loadMoreNode);
+    }
+
+    database.collections.subscribe((collections: ViewModels.Collection[]) => {
+      collections.forEach((collection: ViewModels.Collection) =>
+        databaseNode.children.push(
+          buildCollectionNode(database, collection, isNotebookEnabled, container, refreshActiveTab)
+        )
+      );
+    });
+
+    return databaseNode;
+  });
+
+  return databaseTreeNodes;
+};