ref(cbrs): introduce RegisteredClass into ConfigurableComponent (#7379)

Retrieve configurable component by doing

## option 1

`ConfigurableComponent.get_component_class(namespace).get_from_name(class_name)`
ex)
`ConfigurableComponent.get_component_class("AllocationPolicy").get_from_name("BytesScannedWindowPolicy")`

## option 2
`AllocationPolicy.get_from_name(class_name)`
ex) `AllocationPolicy.get_from_name("BytesScannedWindowPolicy)`


Option 1 is used for the set configuration endpoint - the frontend code
will pass in the namespace and the class name, so use option 1 to get
the configurable component

Option 2 is already used in existing code so I'm going to preserve that


Note: 
When you print out `ConfigurableComponent.all_names()`, you get 

```
['AllocationPolicy.AllocationPolicy', 'AllocationPolicy.PassthroughPolicy', 'AllocationPolicy.BaseConcurrentRateLimitAllocationPolicy', 'AllocationPolicy.ConcurrentRateLimitAllocationPolicy', 'AllocationPolicy.DeleteConcurrentRateLimitAllocationPolicy', 'AllocationPolicy.ReferrerGuardRailPolicy', 'AllocationPolicy.CrossOrgQueryAllocationPolicy', 'AllocationPolicy.BytesScannedRejectingPolicy', 'AllocationPolicy.BytesScannedWindowAllocationPolicy', 'RoutingStrategy.BaseRoutingStrategy', 'RoutingStrategy.OutcomesBasedRoutingStrategy']
```


At one point I did think about doing
`ConfigurableComponent.get_from_name(namespace, class_name)`, similar to
[RPCEndpoint.get_from_name(name,
version)](https://github.com/getsentry/snuba/blob/master/snuba/web/rpc/__init__.py#L133-L137).
But I decided against it because
we either change existing `AllocationPolicy.get_from_name(blah)` to
`ConfigurableComponent.get_from_name("AllocationPolicy", "blah")`
throughout the code base, or within each subclass's `get_from_name`
method, we call `super().get_from_name(subclass's namespace, name)`, and
then we cast it to type subclass, and we make namespace an optional
param. Maybe this works, but I already coded up the existing approach
and I don't see a clear advantage here

---------

Co-authored-by: Rachel Chen <[email protected]>
Co-authored-by: getsantry[bot] <66042841+getsantry[bot]@users.noreply.github.com>

Author: 3 people

snuba/admin/views.py

@@ -175,9 +175,7 @@ def settings_endpoint() -> Response:
 
 @application.route("/tools")
 def tools() -> Response:
-    return make_response(
-        jsonify({"tools": [t.value for t in get_user_allowed_tools(g.user)]}), 200
-    )
+    return make_response(jsonify({"tools": [t.value for t in get_user_allowed_tools(g.user)]}), 200)
 
 
 @application.route("/migrations/groups")
@@ -225,9 +223,7 @@ def migrations_groups_list(group: str) -> Response:
 )
 @check_tool_perms(tools=[AdminTools.MIGRATIONS])
 def run_migration(group: str, migration_id: str) -> Response:
-    return run_or_reverse_migration(
-        group=group, action="run", migration_id=migration_id
-    )
+    return run_or_reverse_migration(group=group, action="run", migration_id=migration_id)
 
 
 @application.route(
@@ -236,19 +232,15 @@ def run_migration(group: str, migration_id: str) -> Response:
 )
 @check_tool_perms(tools=[AdminTools.MIGRATIONS])
 def reverse_migration(group: str, migration_id: str) -> Response:
-    return run_or_reverse_migration(
-        group=group, action="reverse", migration_id=migration_id
-    )
+    return run_or_reverse_migration(group=group, action="reverse", migration_id=migration_id)
 
 
 @application.route(
     "/migrations/<group>/overwrite/<migration_id>/status/<new_status>",
     methods=["POST"],
 )
 @check_tool_perms(tools=[AdminTools.MIGRATIONS])
-def force_overwrite_migration_status(
-    group: str, migration_id: str, new_status: str
-) -> Response:
+def force_overwrite_migration_status(group: str, migration_id: str, new_status: str) -> Response:
     try:
         migration_group = MigrationGroup(group)
     except ValueError as err:
@@ -302,9 +294,7 @@ def do_action() -> None:
         if action == "run":
             runner.run_migration(migration_key, force=force, fake=fake, dry_run=dry_run)
         else:
-            runner.reverse_migration(
-                migration_key, force=force, fake=fake, dry_run=dry_run
-            )
+            runner.reverse_migration(migration_key, force=force, fake=fake, dry_run=dry_run)
 
         if not dry_run:
             audit_log.record(
@@ -356,15 +346,11 @@ def write(self, s: str) -> int:
     except ClickhouseError as err:
         notify_error()
         logger.error(err, exc_info=True)
-        return make_response(
-            jsonify({"error": "clickhouse error: " + err.message}), 400
-        )
+        return make_response(jsonify({"error": "clickhouse error: " + err.message}), 400)
     except InactiveClickhouseReplica as err:
         notify_error()
         logger.error(err, exc_info=True)
-        return make_response(
-            jsonify({"error": "inactive replicas error: " + err.message}), 400
-        )
+        return make_response(jsonify({"error": "inactive replicas error: " + err.message}), 400)
 
 
 @application.route("/clickhouse_queries")
@@ -432,9 +418,7 @@ def clickhouse_system_query() -> Response:
         return make_response(jsonify({"error": "Invalid request"}), 400)
 
     try:
-        result = run_system_query_on_host_with_sql(
-            host, port, storage, raw_sql, sudo_mode, g.user
-        )
+        result = run_system_query_on_host_with_sql(host, port, storage, raw_sql, sudo_mode, g.user)
         rows = []
         rows, columns = cast(List[List[str]], result.results), result.meta
 
@@ -600,16 +584,12 @@ def summarize_trace_with_profile() -> Response:
     except ClickhouseError as err:
         logger.error(err, exc_info=True)
         return make_response(
-            jsonify(
-                {"error": {"type": "clickhouse", "message": str(err), "code": err.code}}
-            ),
+            jsonify({"error": {"type": "clickhouse", "message": str(err), "code": err.code}}),
             400,
         )
     except Exception as err:
         logger.error(err, exc_info=True)
-        return make_response(
-            jsonify({"error": {"type": "unknown", "message": str(err)}}), 500
-        )
+        return make_response(jsonify({"error": {"type": "unknown", "message": str(err)}}), 500)
 
 
 @application.route("/clickhouse_querylog_query", methods=["POST"])
@@ -910,18 +890,14 @@ def serialize(
 @application.route("/clickhouse_nodes")
 @check_tool_perms(tools=[AdminTools.SYSTEM_QUERIES, AdminTools.QUERY_TRACING])
 def clickhouse_nodes() -> Response:
-    return Response(
-        json.dumps(get_storage_info()), 200, {"Content-Type": "application/json"}
-    )
+    return Response(json.dumps(get_storage_info()), 200, {"Content-Type": "application/json"})
 
 
 @application.route("/snuba_datasets")
 @check_tool_perms(tools=[AdminTools.SNQL_TO_SQL])
 def snuba_datasets() -> Response:
     return Response(
-        json.dumps(
-            get_enabled_dataset_names(), 200, {"Content-Type": "application/json"}
-        )
+        json.dumps(get_enabled_dataset_names(), 200, {"Content-Type": "application/json"})
     )
 
 
@@ -972,9 +948,7 @@ def storages_with_allocation_policies() -> Response:
 @check_tool_perms(tools=[AdminTools.CAPACITY_MANAGEMENT])
 def get_allocation_policy_configs(storage_key: str) -> Response:
     storage = get_storage(StorageKey(storage_key))
-    policies = (
-        storage.get_allocation_policies() + storage.get_delete_allocation_policies()
-    )
+    policies = storage.get_allocation_policies() + storage.get_delete_allocation_policies()
 
     return Response(
         json.dumps([convert(policy.to_dict()) for policy in policies]),
@@ -1005,7 +979,7 @@ def set_allocation_policy_config() -> Response:
             + get_storage(StorageKey(storage)).get_delete_allocation_policies()
         )
         policy = next(
-            (p for p in policies if p.config_key() == policy_name),
+            (p for p in policies if p.class_name() == policy_name),
             None,
         )
         assert policy is not None, "Policy not found on storage"
@@ -1022,23 +996,21 @@ def set_allocation_policy_config() -> Response:
         audit_log.record(
             user or "",
             AuditLogAction.ALLOCATION_POLICY_DELETE,
-            {"storage": storage, "policy": policy.config_key(), "key": key},
+            {"storage": storage, "policy": policy.class_name(), "key": key},
             notify=True,
         )
         return Response("", 200)
     elif request.method == "POST":
         try:
             value = data["value"]
             assert isinstance(value, str), "Invalid value"
-            policy.set_config_value(
-                config_key=key, value=value, params=params, user=user
-            )
+            policy.set_config_value(config_key=key, value=value, params=params, user=user)
             audit_log.record(
                 user or "",
                 AuditLogAction.ALLOCATION_POLICY_UPDATE,
                 {
                     "storage": storage,
-                    "policy": policy.config_key(),
+                    "policy": policy.class_name(),
                     "key": key,
                     "value": value,
                     "params": str(params),
@@ -1192,9 +1164,7 @@ def execute_rpc_endpoint(endpoint_name: str, version: str) -> Response:
         )
     except InvalidConfigKeyError:
         return Response(
-            json.dumps(
-                {"error": f"Unknown endpoint: {endpoint_name} or version: {version}"}
-            ),
+            json.dumps({"error": f"Unknown endpoint: {endpoint_name} or version: {version}"}),
             404,
             {"Content-Type": "application/json"},
         )
@@ -1204,9 +1174,7 @@ def execute_rpc_endpoint(endpoint_name: str, version: str) -> Response:
     try:
         request_proto = Parse(json.dumps(body), endpoint_class.request_class()())
         validate_request_meta(request_proto)
-        response = run_rpc_handler(
-            endpoint_name, version, request_proto.SerializeToString()
-        )
+        response = run_rpc_handler(endpoint_name, version, request_proto.SerializeToString())
         return Response(
             json.dumps(MessageToDict(response)),
             200,
@@ -1333,9 +1301,7 @@ def delete() -> Response:
             sentry_sdk.capture_exception(e)
             return make_response(jsonify({"error": "unexpected internal error"}), 500)
 
-    return Response(
-        json.dumps(delete_results), 200, {"Content-Type": "application/json"}
-    )
+    return Response(json.dumps(delete_results), 200, {"Content-Type": "application/json"})
 
 
 @application.route(
@@ -1396,9 +1362,7 @@ def clickhouse_system_settings() -> Response:
     port = request.args.get("port")
     storage = request.args.get("storage")
     if not all([host, port, storage]):
-        return make_response(
-            jsonify({"error": "Host, port, and storage are required"}), 400
-        )
+        return make_response(jsonify({"error": "Host, port, and storage are required"}), 400)
     try:
         # conversions for typing
         settings = get_system_settings(str(host), int(str(port)), str(storage))

snuba/configs/configuration.py

@@ -1,16 +1,19 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field, replace
-from typing import Any, TypedDict, final
+from typing import Any, Type, TypedDict, TypeVar, cast, final
 
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.state import delete_config as delete_runtime_config
 from snuba.state import get_all_configs as get_all_runtime_configs
 from snuba.state import get_config as get_runtime_config
 from snuba.state import set_config as set_runtime_config
+from snuba.utils.registered_class import RegisteredClass
 
 logger = logging.getLogger("snuba.configurable_component")
 
+T = TypeVar("T", bound="ConfigurableComponent")
+
 
 class InvalidConfig(Exception):
     pass
@@ -60,7 +63,7 @@ class Configuration:
     param_types: dict[str, type] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
-        if type(self.default) != self.value_type:
+        if type(self.default) is not self.value_type:
             raise ValueError(
                 f"Config item `{self.name}` expects type {self.value_type} got value `{self.default}` of type {type(self.default)}"
             )
@@ -83,9 +86,7 @@ def to_definition_dict(self) -> dict[str, Any]:
             ],
         }
 
-    def to_config_dict(
-        self, value: Any = None, params: dict[str, Any] = {}
-    ) -> dict[str, Any]:
+    def to_config_dict(self, value: Any = None, params: dict[str, Any] = {}) -> dict[str, Any]:
         """Returns a dict representation of a live Config."""
         return {
             **self.__to_base_dict(),
@@ -94,7 +95,7 @@ def to_config_dict(
         }
 
 
-class ConfigurableComponent(ABC):
+class ConfigurableComponent(ABC, metaclass=RegisteredClass):
     """
     A ConfigurableComponent is a component that can be configured via configurations.
     example: an allocation policy, a routing strategy, a strategy selector.
@@ -148,7 +149,8 @@ def component_name(self) -> str:
         # needs to uniquely identify the configurable component
         return f"{self.resource_identifier.value}.{self.__class__.__name__}"
 
-    def component_namespace(self) -> str:
+    @classmethod
+    def component_namespace(cls) -> str:
         # is it an allocation policy? a routing strategy? a strategy selector?
         raise NotImplementedError
 
@@ -191,9 +193,7 @@ def _validate_config_params(
 
         # config doesn't exist
         if config_key not in definitions:
-            raise InvalidConfig(
-                f"'{config_key}' is not a valid config for {class_name}!"
-            )
+            raise InvalidConfig(f"'{config_key}' is not a valid config for {class_name}!")
 
         config = definitions[config_key]
 
@@ -335,9 +335,7 @@ def _get_overridden_additional_config_defaults(
         return [
             replace(
                 definition,
-                default=default_config_overrides.get(
-                    definition.name, definition.default
-                ),
+                default=default_config_overrides.get(definition.name, definition.default),
             )
             for definition in definitions
         ]
@@ -350,9 +348,7 @@ def __escape_delimiter_chars(self, key: str) -> str:
             escape_sequence,
         ) in self._KEY_DELIMITERS_TO_ESCAPE_SEQUENCES.items():
             if escape_sequence in str(key):
-                raise InvalidConfig(
-                    f"{escape_sequence} is not a valid string for a policy config"
-                )
+                raise InvalidConfig(f"{escape_sequence} is not a valid string for a policy config")
             key = key.replace(delimiter_char, escape_sequence)
         return key
 
@@ -431,13 +427,28 @@ def delete_config_value(
 
     @classmethod
     def config_key(cls) -> str:
+        return f"{cls.component_namespace()}.{cls.__name__}"
+
+    @classmethod
+    def class_name(cls) -> str:
         return cls.__name__
 
     def to_dict(self) -> ConfigurableComponentData:
         return ConfigurableComponentData(
             configurable_component_namespace=self.component_namespace(),
-            configurable_component_config_key=self.config_key(),
+            configurable_component_config_key=self.class_name(),
             resource_identifier=self.resource_identifier.value,
             configurations=self.get_current_configs(),
             optional_config_definitions=self.get_optional_config_definitions_json(),
         )
+
+    @classmethod
+    def get_component_class(cls, namespace: str) -> Type["ConfigurableComponent"]:
+        return cast(
+            Type["ConfigurableComponent"],
+            cls.class_from_name(f"{namespace}.{namespace}"),
+        )
+
+    @classmethod
+    def get_from_name(cls: Type[T], name: str) -> Type[T]:
+        return cast(Type[T], cls.class_from_name(f"{cls.component_namespace()}.{name}"))

snuba/query/allocation_policies/__init__.py

@@ -16,7 +16,7 @@
 )
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.utils.metrics.wrapper import MetricsWrapper
-from snuba.utils.registered_class import RegisteredClass, import_submodules_in_directory
+from snuba.utils.registered_class import import_submodules_in_directory
 from snuba.utils.serializable_exception import JsonSerializable, SerializableException
 from snuba.web import QueryException, QueryResult
 
@@ -127,9 +127,7 @@ def violations(self) -> dict[str, dict[str, Any]]:
 
     @property
     def quota_allowance(self) -> dict[str, dict[str, Any]]:
-        return cast(
-            dict[str, dict[str, Any]], self.extra_data.get("quota_allowances", {})
-        )
+        return cast(dict[str, dict[str, Any]], self.extra_data.get("quota_allowances", {}))
 
     @property
     def summary(self) -> dict[str, Any]:
@@ -155,7 +153,7 @@ class QueryType(Enum):
     DELETE = "delete"
 
 
-class AllocationPolicy(ConfigurableComponent, ABC, metaclass=RegisteredClass):
+class AllocationPolicy(ConfigurableComponent, ABC):
     """This class should be the centralized place for policy decisions regarding
     resource usage of a clickhouse cluster. It is meant to live as a configurable item
     on a storage.
@@ -373,7 +371,8 @@ def __init__(
             self._get_overridden_additional_config_defaults(default_config_overrides)
         )
 
-    def component_namespace(self) -> str:
+    @classmethod
+    def component_namespace(cls) -> str:
         return "AllocationPolicy"
 
     def _get_hash(self) -> str:
@@ -393,10 +392,7 @@ def metrics(self) -> MetricsWrapper:
 
     @property
     def is_active(self) -> bool:
-        return (
-            bool(self.get_config_value(IS_ACTIVE))
-            and settings.ALLOCATION_POLICY_ENABLED
-        )
+        return bool(self.get_config_value(IS_ACTIVE)) and settings.ALLOCATION_POLICY_ENABLED
 
     @property
     def is_enforced(self) -> bool:
@@ -407,10 +403,6 @@ def max_threads(self) -> int:
         """Maximum number of threads run a single query on ClickHouse with."""
         return int(self.get_config_value(MAX_THREADS))
 
-    @classmethod
-    def get_from_name(cls, name: str) -> "AllocationPolicy":
-        return cast("AllocationPolicy", cls.class_from_name(name))
-
     def __eq__(self, other: Any) -> bool:
         """There should not be a need to compare these except that
         AllocationPolicies are attached to the Table a query is executed against.