feat(cbrs): UI

Author: Rachel Chen

snuba/admin/rpc/rpc_queries.py

@@ -2,6 +2,9 @@
 
 from snuba import settings
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
+from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
+    BaseRoutingStrategy,
+)
 
 
 def _validate_projects_in_query(project_ids: List[int]) -> None:
@@ -30,3 +33,7 @@ def validate_request_meta(request_proto: Any) -> None:
     project_ids = list(meta.project_ids)
     _validate_org_ids_in_query(org_id)
     _validate_projects_in_query(project_ids)
+
+
+def get_routing_strategies() -> List[str]:
+    return list(BaseRoutingStrategy.all_names())

snuba/admin/static/api_client.tsx

@@ -35,7 +35,7 @@ import {
   CardinalityQueryResult,
 } from "SnubaAdmin/cardinality_analyzer/types";
 
-import { AllocationPolicy } from "SnubaAdmin/capacity_management/types";
+import { AllocationPolicy, Entity } from "SnubaAdmin/capacity_management/types";
 
 import { ReplayInstruction, Topic } from "SnubaAdmin/dead_letter_queue/types";
 import { AutoReplacementsBypassProjectsData } from "SnubaAdmin/auto_replacements_bypass_projects/types";
@@ -83,17 +83,18 @@ interface Client {
   getAllMigrationGroups: () => Promise<MigrationGroupResult[]>;
   runMigration: (req: RunMigrationRequest) => Promise<RunMigrationResult>;
   getAllowedTools: () => Promise<AllowedTools>;
+  getRoutingStrategies: () => Promise<string[]>;
   getStoragesWithAllocationPolicies: () => Promise<string[]>;
-  getAllocationPolicies: (storage: string) => Promise<AllocationPolicy[]>;
+  getAllocationPolicies: (entity: Entity) => Promise<AllocationPolicy[]>;
   setAllocationPolicyConfig: (
-    storage: string,
+    entity: Entity,
     policy: string,
     key: string,
     value: string,
     params: object,
   ) => Promise<void>;
   deleteAllocationPolicyConfig: (
-    storage: string,
+    entity_name: string,
     policy: string,
     key: string,
     params: object,
@@ -445,31 +446,39 @@ function Client(): Client {
       }).then((resp) => resp.json());
     },
 
+    getRoutingStrategies: () => {
+      const url = baseUrl + "routing_strategies";
+      return fetch(url, {
+        headers: { "Content-Type": "application/json" },
+      }).then((resp) => resp.json());
+    },
+
     getStoragesWithAllocationPolicies: () => {
       const url = baseUrl + "storages_with_allocation_policies";
       return fetch(url, {
         headers: { "Content-Type": "application/json" },
       }).then((resp) => resp.json());
     },
-    getAllocationPolicies: (storage: string) => {
+    getAllocationPolicies: (entity: Entity) => {
       const url =
-        baseUrl + "allocation_policy_configs/" + encodeURIComponent(storage);
+        baseUrl + "allocation_policy_configs/" + entity.type + "/" + encodeURIComponent(entity.name);
       return fetch(url, {
         headers: { "Content-Type": "application/json" },
       }).then((resp) => resp.json());
     },
     setAllocationPolicyConfig: (
-      storage: string,
+      entity: Entity,
       policy: string,
       key: string,
       value: string,
       params: object,
     ) => {
+      const body = entity.type === "storage" ? JSON.stringify({ storage: entity.name, policy, key, value, params }) : JSON.stringify({ strategy: entity.name, policy, key, value, params });
       const url = baseUrl + "allocation_policy_config";
       return fetch(url, {
         headers: { "Content-Type": "application/json" },
         method: "POST",
-        body: JSON.stringify({ storage, policy, key, value, params }),
+        body: body,
       }).then((res) => {
         if (res.ok) {
           return;
@@ -482,7 +491,7 @@ function Client(): Client {
       });
     },
     deleteAllocationPolicyConfig: (
-      storage: string,
+      entity_name: string,
       policy: string,
       key: string,
       params: object,
@@ -491,7 +500,7 @@ function Client(): Client {
       return fetch(url, {
         headers: { "Content-Type": "application/json" },
         method: "DELETE",
-        body: JSON.stringify({ storage, policy, key, params }),
+        body: JSON.stringify({ entity_name, policy, key, params }),
       }).then((res) => {
         if (res.ok) {
           return;

snuba/admin/static/capacity_management/allocation_policy.tsx

@@ -3,7 +3,7 @@ import React, { useEffect, useState } from "react";
 import { Table, createCustomTableStyles } from "../table";
 import { COLORS } from "SnubaAdmin/theme";
 import Client from "SnubaAdmin/api_client";
-import { AllocationPolicy, AllocationPolicyConfig } from "SnubaAdmin/capacity_management/types";
+import { AllocationPolicy, AllocationPolicyConfig, Entity } from "SnubaAdmin/capacity_management/types";
 import { containerStyle, linkStyle, paragraphStyle } from "SnubaAdmin/capacity_management/styles";
 import { getReadonlyRow } from "SnubaAdmin/capacity_management/row_data";
 import EditConfigModal from "SnubaAdmin/capacity_management/edit_config_modal";
@@ -39,10 +39,10 @@ function getTableColor(configs: AllocationPolicyConfig[]): string {
 
 function AllocationPolicyConfigs(props: {
   api: Client;
-  storage: string;
+  entity: Entity;
   policy: AllocationPolicy;
 }) {
-  const { api, storage, policy } = props;
+  const { api, entity, policy } = props;
 
   const [configs, setConfigs] = useState<AllocationPolicyConfig[]>([]);
 
@@ -69,7 +69,7 @@ function AllocationPolicyConfigs(props: {
   function deleteConfig(toDelete: AllocationPolicyConfig) {
     api
       .deleteAllocationPolicyConfig(
-        storage,
+        entity.name,
         policy.policy_name,
         toDelete.name,
         toDelete.params
@@ -89,7 +89,7 @@ function AllocationPolicyConfigs(props: {
   function saveConfig(config: AllocationPolicyConfig) {
     api
       .setAllocationPolicyConfig(
-        storage,
+        entity,
         policy.policy_name,
         config.name,
         config.value,

snuba/admin/static/capacity_management/index.tsx

@@ -31,7 +31,7 @@ function CapacityManagement(props: { api: Client }) {
 
   function loadAllocationPolicies(storage: string) {
     api
-      .getAllocationPolicies(storage)
+      .getAllocationPolicies({ type: "storage", name: storage })
       .then((res) => {
         setAllocationPolicies(res);
       })
@@ -55,7 +55,7 @@ function CapacityManagement(props: { api: Client }) {
         {policies.map((policy: AllocationPolicy) => (
           <AllocationPolicyConfigs
             api={api}
-            storage={selectedStorage}
+            entity={{ type: "storage", name: selectedStorage }}
             policy={policy}
             key={selectedStorage + policy.policy_name}
           />

snuba/admin/static/capacity_management/types.tsx

@@ -37,10 +37,41 @@ type RowData = {
   edit: ReactNode;
 };
 
+interface StorageEntity {
+  type: "storage";
+  name: string;
+}
+
+interface StrategyEntity {
+  type: "strategy";
+  name: string;
+}
+
+type Entity = StorageEntity | StrategyEntity;
+
+// Usage
+function getEntityName(entity: Entity): string {
+  return entity.name; // Works for both types
+}
+
+function isStorage(entity: Entity): entity is StorageEntity {
+  return entity.type === "storage";
+}
+
+function isStrategy(entity: Entity): entity is StrategyEntity {
+  return entity.type === "strategy";
+}
+
 export {
   AllocationPolicy,
   AllocationPolicyConfig,
   AllocationPolicyOptionalConfigDefinition,
   AllocationPolicyConfigParams,
   RowData,
+  Entity,
+  getEntityName,
+  isStorage,
+  isStrategy,
+  StorageEntity,
+  StrategyEntity,
 };

snuba/admin/static/cbrs/index.tsx

@@ -0,0 +1,99 @@
+import React, { useEffect, useState } from "react";
+import Client from "SnubaAdmin/api_client";
+import { AllocationPolicyConfigs } from "SnubaAdmin/capacity_management/allocation_policy";
+import { AllocationPolicy } from "SnubaAdmin/capacity_management/types";
+import { CustomSelect, getParamFromStorage } from "SnubaAdmin/select";
+import { COLORS } from "SnubaAdmin/theme";
+
+function CapacityBasedRoutingSystem(props: { api: Client }) {
+  const { api } = props;
+
+  const [strategies, setStrategies] = useState<string[]>([]);
+  const [selectedStrategy, setStrategy] = useState<string | undefined>();
+  const [allocationPolicies, setAllocationPolicies] = useState<
+    AllocationPolicy[]
+  >([]);
+
+  useEffect(() => {
+    api.getRoutingStrategies().then((res) => {
+      setStrategies(res);
+      const previousStrategy = getParamFromStorage("strategy");
+      if (previousStrategy) {
+        selectStrategy(previousStrategy);
+      }
+    });
+  }, []);
+
+  function selectStrategy(strategy: string) {
+    setStrategy(strategy);
+    loadAllocationPolicies(strategy);
+  }
+
+  function loadAllocationPolicies(strategy: string) {
+    api
+      .getAllocationPolicies({ type: "strategy", name: strategy })
+      .then((res) => {
+        setAllocationPolicies(res);
+      })
+      .catch((err) => {
+        window.alert(err);
+      });
+  }
+
+  function renderPolicies(policies: AllocationPolicy[]) {
+    if (!selectedStrategy) {
+      return <p>Strategy not selected.</p>;
+    }
+    if (policies.length == 0) {
+      return null;
+    }
+    return (
+      <div>
+        <p style={policyTypeStyle}>
+          Policy Type: {policies[0].query_type.toUpperCase()}
+        </p>
+        {policies.map((policy: AllocationPolicy) => (
+          <AllocationPolicyConfigs
+            api={api}
+            entity={{ type: "strategy", name: selectedStrategy }}
+            policy={policy}
+            key={selectedStrategy + policy.policy_name}
+          />
+        ))}
+      </div>
+    );
+  }
+
+  return (
+    <div>
+      <p>
+        Strategy:
+        <CustomSelect
+          value={selectedStrategy || ""}
+          onChange={selectStrategy}
+          name="strategy"
+          options={strategies}
+        />
+      </p>
+
+      {renderPolicies(
+        allocationPolicies.filter((policy) => policy.query_type == "select")
+      )}
+      {renderPolicies(
+        allocationPolicies.filter((policy) => policy.query_type == "delete")
+      )}
+    </div>
+  );
+}
+
+const policyTypeStyle = {
+  fontSize: 18,
+  fontWeight: 600,
+  color: COLORS.HEADER_TEXT,
+  backgroundColor: COLORS.TEXT_LIGHTER,
+  maxWidth: "100%",
+  margin: "10px 0px",
+  padding: "5px",
+};
+
+export default CapacityBasedRoutingSystem;

snuba/admin/static/data.tsx

@@ -18,6 +18,7 @@ import DeleteTool from "SnubaAdmin/delete_tool";
 import ViewCustomJobs from "SnubaAdmin/manual_jobs";
 import DatabaseClusters from "./database_clusters";
 import RpcEndpoints from "SnubaAdmin/rpc_endpoints";
+import CapacityBasedRoutingSystem from "SnubaAdmin/cbrs";
 
 const NAV_ITEMS = [
   { id: "overview", display: "🤿 Snuba Admin", component: Welcome },
@@ -112,6 +113,11 @@ const NAV_ITEMS = [
     display: "🗂️ Database Clusters",
     component: DatabaseClusters,
   },
+  {
+    id: "capacity-based-routing-system",
+    display: "🔄 Capacity Based Routing System",
+    component: CapacityBasedRoutingSystem,
+  },
 ];
 
 export { NAV_ITEMS };

snuba/admin/tool_policies.py

@@ -28,6 +28,7 @@ class AdminTools(Enum):
     AUDIT_LOG = "audit-log"
     KAFKA = "kafka"
     CAPACITY_MANAGEMENT = "capacity-management"
+    CAPACITY_BASED_ROUTING_SYSTEM = "capacity-based-routing-system"
     PRODUCTION_QUERIES = "production-queries"
     CARDINALITY_ANALYZER = "cardinality-analyzer"
     SNUBA_EXPLAIN = "snuba-explain"