cosmos-explorer/src/Explorer/Tree/ResourceTree.tsx

import { Tree, TreeItemValue, TreeOpenChangeData, TreeOpenChangeEvent } from "@fluentui/react-components";
import { Home16Regular } from "@fluentui/react-icons";
import { AuthType } from "AuthType";
import { useTreeStyles } from "Explorer/Controls/TreeComponent/Styles";
import { TreeNode, TreeNodeComponent } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
import {
  createDatabaseTreeNodes,
  createResourceTokenTreeNodes,
  createSampleDataTreeNodes,
} from "Explorer/Tree/treeNodeUtil";
import { useDatabases } from "Explorer/useDatabases";
import { useSelectedNode } from "Explorer/useSelectedNode";
import { isFabricMirrored } from "Platform/Fabric/FabricUtil";
import { userContext } from "UserContext";
import { useQueryCopilot } from "hooks/useQueryCopilot";
import { ReactTabKind, useTabs } from "hooks/useTabs";
import * as React from "react";
import { useEffect, useMemo } 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 {
  explorer: Explorer;
}

export const DATA_TREE_LABEL = "DATA";
export const MY_DATA_TREE_LABEL = "MY DATA";
export const SAMPLE_DATA_TREE_LABEL = "SAMPLE DATA";

/**
 * Top-level tree that has no label, but contains all subtrees
 */
export const ResourceTree: React.FC<ResourceTreeProps> = ({ explorer }: ResourceTreeProps): JSX.Element => {
  const [openItems, setOpenItems] = React.useState<TreeItemValue[]>([]);
  const treeStyles = useTreeStyles();

  const { isNotebookEnabled } = useNotebook(
    (state) => ({
      isNotebookEnabled: state.isNotebookEnabled,
    }),
    shallow,
  );

  // We intentionally avoid using a state selector here because we want to re-render the tree if the active tab changes.
  const { refreshActiveTab } = useTabs();

  const { databases, resourceTokenCollection, sampleDataResourceTokenCollection } = useDatabases((state) => ({
    databases: state.databases,
    resourceTokenCollection: state.resourceTokenCollection,
    sampleDataResourceTokenCollection: state.sampleDataResourceTokenCollection,
  }));
  const { isCopilotEnabled, isCopilotSampleDBEnabled } = useQueryCopilot((state) => ({
    isCopilotEnabled: state.copilotEnabled,
    isCopilotSampleDBEnabled: state.copilotSampleDBEnabled,
  }));

  const databaseTreeNodes = useMemo(() => {
    return userContext.authType === AuthType.ResourceToken
      ? createResourceTokenTreeNodes(resourceTokenCollection)
      : createDatabaseTreeNodes(explorer, isNotebookEnabled, databases, refreshActiveTab);
  }, [resourceTokenCollection, databases, isNotebookEnabled, refreshActiveTab]);

  const isSampleDataEnabled =
    isCopilotEnabled &&
    isCopilotSampleDBEnabled &&
    userContext.sampleDataConnectionInfo &&
    userContext.apiType === "SQL";

  const sampleDataNodes = useMemo<TreeNode[]>(() => {
    return isSampleDataEnabled && sampleDataResourceTokenCollection
      ? createSampleDataTreeNodes(sampleDataResourceTokenCollection)
      : [];
  }, [isSampleDataEnabled, sampleDataResourceTokenCollection]);

  const headerNodes: TreeNode[] = isFabricMirrored()
    ? []
    : [
        {
          id: "home",
          iconSrc: <Home16Regular />,
          label: "Home",
          isSelected: () =>
            useSelectedNode.getState().selectedNode === undefined &&
            useTabs.getState().activeReactTab === ReactTabKind.Home,
          onClick: () => {
            useSelectedNode.getState().setSelectedNode(undefined);
            useTabs.getState().openAndActivateReactTab(ReactTabKind.Home);
          },
        },
      ];

  const rootNodes: TreeNode[] = useMemo(() => {
    if (sampleDataNodes.length > 0) {
      return [
        ...headerNodes,
        {
          id: "data",
          label: MY_DATA_TREE_LABEL,
          children: databaseTreeNodes,
          isScrollable: true,
        },
        {
          id: "sampleData",
          label: SAMPLE_DATA_TREE_LABEL,
          children: sampleDataNodes,
        },
      ];
    } else {
      return [...headerNodes, ...databaseTreeNodes];
    }
  }, [databaseTreeNodes, sampleDataNodes]);

  useEffect(() => {
    // Compute open items based on node.isExpanded
    const updateOpenItems = (node: TreeNode, parentNodeId: string): void => {
      // This will look for ANY expanded node, event if its parent node isn't expanded
      // and add it to the openItems list
      const globalId = parentNodeId === undefined ? node.label : `${parentNodeId}/${node.label}`;

      if (node.isExpanded) {
        let found = false;
        for (const id of openItems) {
          if (id === globalId) {
            found = true;
            break;
          }
        }
        if (!found) {
          setOpenItems((prevOpenItems) => [...prevOpenItems, globalId]);
        }
      }

      if (node.children) {
        for (const child of node.children) {
          updateOpenItems(child, globalId);
        }
      }
    };

    rootNodes.forEach((n) => updateOpenItems(n, undefined));
  }, [rootNodes, openItems, setOpenItems]);

  const handleOpenChange = (event: TreeOpenChangeEvent, data: TreeOpenChangeData) =>
    setOpenItems(Array.from(data.openItems));

  return (
    <div className={treeStyles.treeContainer}>
      <Tree
        aria-label="CosmosDB resources"
        openItems={openItems}
        className={treeStyles.tree}
        onOpenChange={handleOpenChange}
        size="medium"
      >
        {rootNodes.map((node) => (
          <TreeNodeComponent
            key={node.label}
            openItems={openItems}
            className="dataResourceTree"
            node={node}
            treeNodeId={node.label}
          />
        ))}
      </Tree>
    </div>
  );
};