ui: cleanup useNodeStatuses usages

This commit cleans up the useNodeStatuses hook. The raw data from the response is no longer returned from the hook. Instead, callers must use the transformed data. Epic: none Release note: None
cockroachdb · Oct 29, 2024 · c09ef5d · c09ef5d
1 parent 1783178
commit c09ef5d
Show file tree

Hide file tree

Showing 10 changed files with 86 additions and 63 deletions.
diff --git a/pkg/ui/workspaces/cluster-ui/src/api/nodesApi.ts b/pkg/ui/workspaces/cluster-ui/src/api/nodesApi.ts
@@ -20,6 +20,11 @@ export const getNodes =
     return fetchData(cockroach.server.serverpb.NodesResponse, NODES_PATH);
   };
 
+export type NodeStatus = {
+  region: string;
+  stores: StoreID[];
+};
+
 export const useNodeStatuses = () => {
   const clusterDetails = useContext(ClusterDetailsContext);
   const isTenant = clusterDetails.isTenant;
@@ -32,29 +37,31 @@ export const useNodeStatuses = () => {
     },
   );
 
-  const { storeIDToNodeID, nodeIDToRegion } = useMemo(() => {
-    const nodeIDToRegion: Record<NodeID, string> = {};
+  const { storeIDToNodeID, nodeStatusByID } = useMemo(() => {
+    const nodeStatusByID: Record<NodeID, NodeStatus> = {};
     const storeIDToNodeID: Record<StoreID, NodeID> = {};
     if (!data) {
-      return { nodeIDToRegion, storeIDToNodeID };
+      return { nodeStatusByID, storeIDToNodeID };
     }
-    data.nodes.forEach(ns => {
-      ns.store_statuses.forEach(store => {
+    data.nodes?.forEach(ns => {
+      ns.store_statuses?.forEach(store => {
         storeIDToNodeID[store.desc.store_id as StoreID] = ns.desc
           .node_id as NodeID;
       });
-      nodeIDToRegion[ns.desc.node_id as NodeID] = getRegionFromLocality(
-        ns.desc.locality,
-      );
+
+      nodeStatusByID[ns.desc.node_id as NodeID] = {
+        region: getRegionFromLocality(ns.desc.locality),
+        stores: ns.store_statuses?.map(s => s.desc.store_id as StoreID),
+      };
     });
-    return { nodeIDToRegion, storeIDToNodeID };
+
+    return { nodeStatusByID, storeIDToNodeID };
   }, [data]);
 
   return {
-    data,
     isLoading,
     error,
-    nodeIDToRegion,
+    nodeStatusByID,
     storeIDToNodeID,
   };
 };
diff --git a/pkg/ui/workspaces/cluster-ui/src/components/nodeRegionsSelector/nodeRegionsSelector.spec.tsx b/pkg/ui/workspaces/cluster-ui/src/components/nodeRegionsSelector/nodeRegionsSelector.spec.tsx
@@ -3,54 +3,54 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
+import { cockroach } from "@cockroachlabs/crdb-protobuf-client";
 import { render, screen, fireEvent, waitFor } from "@testing-library/react";
 import React from "react";
 
-import { useNodeStatuses } from "src/api";
-import { getRegionFromLocality } from "src/store/nodes";
-import { StoreID } from "src/types/clusterTypes";
+import * as api from "src/api/nodesApi";
+import { NodeID, StoreID } from "src/types/clusterTypes";
 
 import { NodeRegionsSelector } from "./nodeRegionsSelector";
 
-// Mock the useNodeStatuses hook
-jest.mock("src/api", () => ({
-  useNodeStatuses: jest.fn(),
-}));
+import NodesResponse = cockroach.server.serverpb.NodesResponse;
 
-// Mock the getRegionFromLocality function
-jest.mock("src/store/nodes", () => ({
-  getRegionFromLocality: jest.fn(),
-}));
-
-const mockNodeData = {
+const mockNodeData = new NodesResponse({
   nodes: [
     {
-      desc: { node_id: 1, locality: { region: "us-east" } },
+      desc: {
+        node_id: 1,
+        locality: { tiers: [{ key: "region", value: "us-east" }] },
+      },
       store_statuses: [
         { desc: { store_id: 101 } },
         { desc: { store_id: 102 } },
       ],
     },
     {
-      desc: { node_id: 2, locality: { region: "us-west" } },
+      desc: {
+        node_id: 2,
+        locality: { tiers: [{ key: "region", value: "us-west" }] },
+      },
       store_statuses: [{ desc: { store_id: 201 } }],
     },
     {
-      desc: { node_id: 3, locality: { region: "us-east" } },
+      desc: {
+        node_id: 3,
+        locality: { tiers: [{ key: "region", value: "us-east" }] },
+      },
       store_statuses: [{ desc: { store_id: 301 } }],
     },
   ],
-};
+});
 
 describe("NodeRegionsSelector", () => {
   beforeEach(() => {
-    (useNodeStatuses as jest.Mock).mockReturnValue({
-      isLoading: false,
-      data: mockNodeData,
-    });
-    (getRegionFromLocality as jest.Mock).mockImplementation(
-      locality => locality.region,
-    );
+    // Mock the api.getNodes function at the module level
+    jest.spyOn(api, "getNodes").mockResolvedValue(mockNodeData);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
   });
 
   it("should render", () => {
@@ -103,9 +103,20 @@ describe("NodeRegionsSelector", () => {
   });
 
   it("handles loading state", () => {
-    (useNodeStatuses as jest.Mock).mockReturnValue({
+    jest.spyOn(api, "useNodeStatuses").mockReturnValue({
+      error: null,
       isLoading: true,
-      data: mockNodeData,
+      nodeStatusByID: {
+        1: { region: "us-east", stores: [101, 102] },
+        2: { region: "us-west", stores: [201] },
+        3: { region: "us-east", stores: [301] },
+      } as Record<NodeID, api.NodeStatus>,
+      storeIDToNodeID: {
+        101: 1,
+        102: 1,
+        201: 2,
+        301: 3,
+      } as Record<StoreID, NodeID>,
     });
 
     render(<NodeRegionsSelector value={[]} onChange={() => {}} />);

diff --git a/pkg/ui/workspaces/cluster-ui/src/components/nodeRegionsSelector/nodeRegionsSelector.tsx b/pkg/ui/workspaces/cluster-ui/src/components/nodeRegionsSelector/nodeRegionsSelector.tsx
@@ -7,7 +7,6 @@ import React, { useMemo } from "react";
 import Select, { OptionsType } from "react-select";
 
 import { useNodeStatuses } from "src/api";
-import { getRegionFromLocality } from "src/store/nodes";
 import { NodeID, StoreID } from "src/types/clusterTypes";
 import {
   GroupedReactSelectOption,
@@ -23,22 +22,24 @@ export const NodeRegionsSelector: React.FC<NodeRegionsSelectorProps> = ({
   value,
   onChange,
 }) => {
-  const nodesResp = useNodeStatuses();
+  const { nodeStatusByID, isLoading } = useNodeStatuses();
 
   const nodeOptions: GroupedReactSelectOption<StoreID[]>[] = useMemo(() => {
     const optionsMap: Record<string, { nid: NodeID; sids: StoreID[] }[]> = {};
-    if (nodesResp.isLoading && !nodesResp.data?.nodes) {
+    const nodes = Object.keys(nodeStatusByID ?? {});
+    if (isLoading && !nodes.length) {
       return [];
     }
 
-    nodesResp.data.nodes.forEach(node => {
-      const region = getRegionFromLocality(node.desc.locality);
+    nodes.forEach(node => {
+      const nid = parseInt(node) as NodeID;
+      const region = nodeStatusByID[nid].region;
       if (optionsMap[region] == null) {
         optionsMap[region] = [];
       }
       optionsMap[region].push({
-        nid: node.desc.node_id as NodeID,
-        sids: node.store_statuses.map(s => s.desc.store_id as StoreID),
+        nid,
+        sids: nodeStatusByID[nid].stores,
       });
     });
 
@@ -51,7 +52,7 @@ export const NodeRegionsSelector: React.FC<NodeRegionsSelectorProps> = ({
         })),
       };
     });
-  }, [nodesResp]);
+  }, [nodeStatusByID, isLoading]);
 
   const onSelectChange = (
     selected: OptionsType<ReactSelectOption<StoreID[]>>,

diff --git a/pkg/ui/workspaces/cluster-ui/src/databaseDetailsV2/tablesView.tsx b/pkg/ui/workspaces/cluster-ui/src/databaseDetailsV2/tablesView.tsx
@@ -240,13 +240,13 @@ export const TablesPageV2 = () => {
   const tableData = useMemo(
     () =>
       tableMetadataToRows(tableList ?? [], {
-        nodeIDToRegion: nodesResp.nodeIDToRegion,
+        nodeStatusByID: nodesResp.nodeStatusByID,
         storeIDToNodeID: nodesResp.storeIDToNodeID,
         isLoading: nodesResp.isLoading,
       }),
     [
       tableList,
-      nodesResp.nodeIDToRegion,
+      nodesResp.nodeStatusByID,
       nodesResp.storeIDToNodeID,
       nodesResp.isLoading,
     ],

diff --git a/pkg/ui/workspaces/cluster-ui/src/databaseDetailsV2/utils.tsx b/pkg/ui/workspaces/cluster-ui/src/databaseDetailsV2/utils.tsx
@@ -3,6 +3,7 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
+import { NodeStatus } from "src/api";
 import { TableMetadata } from "src/api/databases/getTableMetadataApi";
 import { NodeID, StoreID } from "src/types/clusterTypes";
 import { mapStoreIDsToNodeRegions } from "src/util/nodeUtils";
@@ -12,15 +13,15 @@ import { TableRow } from "./types";
 export const tableMetadataToRows = (
   tables: TableMetadata[],
   nodesInfo: {
-    nodeIDToRegion: Record<NodeID, string>;
+    nodeStatusByID: Record<NodeID, NodeStatus>;
     storeIDToNodeID: Record<StoreID, NodeID>;
     isLoading: boolean;
   },
 ): TableRow[] => {
   return tables.map(table => {
     const nodesByRegion = mapStoreIDsToNodeRegions(
       table.storeIds,
-      nodesInfo?.nodeIDToRegion,
+      nodesInfo?.nodeStatusByID,
       nodesInfo?.storeIDToNodeID,
     );
     return {

diff --git a/pkg/ui/workspaces/cluster-ui/src/databasesV2/index.tsx b/pkg/ui/workspaces/cluster-ui/src/databasesV2/index.tsx
@@ -159,7 +159,7 @@ export const DatabasesPageV2 = () => {
   const tableData = useMemo(
     () =>
       rawDatabaseMetadataToDatabaseRows(data?.results ?? [], {
-        nodeIDToRegion: nodesResp.nodeIDToRegion,
+        nodeStatusByID: nodesResp.nodeStatusByID,
         storeIDToNodeID: nodesResp.storeIDToNodeID,
         isLoading: nodesResp.isLoading,
       }),

diff --git a/pkg/ui/workspaces/cluster-ui/src/databasesV2/utils.ts b/pkg/ui/workspaces/cluster-ui/src/databasesV2/utils.ts
@@ -4,6 +4,7 @@
 // included in the /LICENSE file.
 
 import { DatabaseMetadata } from "src/api/databases/getDatabaseMetadataApi";
+import { NodeStatus } from "src/api/nodesApi";
 import { NodeID, StoreID } from "src/types/clusterTypes";
 import { mapStoreIDsToNodeRegions } from "src/util/nodeUtils";
 
@@ -12,15 +13,15 @@ import { DatabaseRow } from "./databaseTypes";
 export const rawDatabaseMetadataToDatabaseRows = (
   raw: DatabaseMetadata[],
   nodesInfo: {
-    nodeIDToRegion: Record<NodeID, string>;
+    nodeStatusByID: Record<NodeID, NodeStatus>;
     storeIDToNodeID: Record<StoreID, NodeID>;
     isLoading: boolean;
   },
 ): DatabaseRow[] => {
   return raw.map((db: DatabaseMetadata): DatabaseRow => {
     const nodesByRegion = mapStoreIDsToNodeRegions(
       db.storeIds,
-      nodesInfo?.nodeIDToRegion,
+      nodesInfo?.nodeStatusByID,
       nodesInfo?.storeIDToNodeID,
     );
     return {

diff --git a/pkg/ui/workspaces/cluster-ui/src/tableDetailsV2/tableOverview.tsx b/pkg/ui/workspaces/cluster-ui/src/tableDetailsV2/tableOverview.tsx
@@ -29,7 +29,7 @@ export const TableOverview: React.FC<TableOverviewProps> = ({
   const isTenant = clusterDetails.isTenant;
   const metadata = tableDetails.metadata;
   const {
-    nodeIDToRegion,
+    nodeStatusByID,
     storeIDToNodeID,
     isLoading: nodesLoading,
   } = useNodeStatuses();
@@ -42,7 +42,7 @@ export const TableOverview: React.FC<TableOverviewProps> = ({
     }
     const regionsToNodes = mapStoreIDsToNodeRegions(
       tableDetails.metadata.storeIds,
-      nodeIDToRegion,
+      nodeStatusByID,
       storeIDToNodeID,
     );
     return Object.entries(regionsToNodes)

diff --git a/pkg/ui/workspaces/cluster-ui/src/util/nodeUtils.spec.ts b/pkg/ui/workspaces/cluster-ui/src/util/nodeUtils.spec.ts
@@ -3,7 +3,8 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-import { StoreID } from "src/types/clusterTypes";
+import { NodeStatus } from "src/api";
+import { NodeID, StoreID } from "src/types/clusterTypes";
 
 import { mapStoreIDsToNodeRegions } from "./nodeUtils";
 
@@ -12,12 +13,12 @@ describe("nodeUtils", () => {
     it("should return a mapping of regions to the nodes that are present in that region based on the provided storeIDs", () => {
       const stores = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] as StoreID[];
       const clusterNodeIDToRegion = {
-        1: "region1",
-        2: "region2",
-        3: "region1",
-        4: "region2",
-        5: "region1",
-      };
+        1: { region: "region1", stores: [1, 6] },
+        2: { region: "region2", stores: [2, 7] },
+        3: { region: "region1", stores: [3, 8] },
+        4: { region: "region2", stores: [4, 9] },
+        5: { region: "region1", stores: [5, 10] },
+      } as Record<NodeID, NodeStatus>;
       const clusterStoreIDToNodeID = {
         1: 1,
         2: 2,

diff --git a/pkg/ui/workspaces/cluster-ui/src/util/nodeUtils.ts b/pkg/ui/workspaces/cluster-ui/src/util/nodeUtils.ts
@@ -3,14 +3,15 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
+import { NodeStatus } from "src/api";
 import { NodeID, StoreID } from "src/types/clusterTypes";
 
 // mapStoreIDsToNodeRegions creates a mapping of regions
 // to the nodes that are present in that region based on
 // the provided storeIDs.
 export const mapStoreIDsToNodeRegions = (
   stores: StoreID[],
-  clusterNodeIDToRegion: Record<NodeID, string> = {},
+  clusterNodeIDToRegion: Record<NodeID, NodeStatus> = {},
   clusterStoreIDToNodeID: Record<StoreID, NodeID> = {},
 ): Record<string, NodeID[]> => {
   const nodes = stores.reduce((acc, storeID) => {
@@ -20,7 +21,7 @@ export const mapStoreIDsToNodeRegions = (
 
   const nodesByRegion: Record<string, NodeID[]> = {};
   nodes.forEach(nodeID => {
-    const region = clusterNodeIDToRegion[nodeID];
+    const region = clusterNodeIDToRegion[nodeID].region;
     if (!nodesByRegion[region]) {
       nodesByRegion[region] = [];
     }