Remove default entity from base and add get_all_entities

Author: enochtangg

bin/mocks/mock-subscriptions

@@ -1,9 +1,9 @@
 #!/usr/bin/env python
 import argparse
 import random
-from datetime import timedelta
 
-from snuba.datasets.entities.factory import get_entity_name
+from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.entity_subscriptions.factory import get_entity_subscription
 from snuba.datasets.factory import get_dataset
 from snuba.redis import redis_client
@@ -31,12 +31,21 @@ parser.add_argument(
     choices=["events"],
     default="events",
 )
+parser.add_argument(
+    "--entity",
+    type=str,
+    help="Entity for subscriptions",
+    dest="entity",
+    choices=["events"],
+    default="events",
+)
 parsed = parser.parse_args()
 
 dataset_name = parsed.dataset
 dataset = get_dataset(dataset_name)
-entity = dataset.get_default_entity()
-entity_key = get_entity_name(entity)
+entity_str = parsed.entity
+entity_key = EntityKey(entity_str)
+entity = get_entity(entity_key)
 storage = entity.get_writable_storage()
 assert storage is not None
 stream_loader = storage.get_table_writer().get_stream_loader()
@@ -52,8 +61,8 @@ for _ in range(parsed.number):
     subscription_data = SubscriptionData(
         query=f"MATCH ({entity_key.value}) SELECT count() AS count WHERE project_id = {project_id}",
         project_id=project_id,
-        time_window=timedelta(minutes=1),
-        resolution=timedelta(minutes=1),
+        time_window_sec=60,
+        resolution_sec=60,
         entity_subscription=entity_subscription,
     )
 

snuba/datasets/cdc/groupassignee.py

@@ -1,5 +1,9 @@
+from typing import Sequence
+
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.entity import Entity
 
 
 class GroupAssigneeDataset(Dataset):
@@ -11,5 +15,7 @@ class GroupAssigneeDataset(Dataset):
     there is no issue in fixing the name.
     """
 
-    def __init__(self) -> None:
-        super().__init__(default_entity=EntityKey.GROUPASSIGNEE)
+    def get_all_entities(self) -> Sequence[Entity]:
+        return [
+            get_entity(EntityKey.GROUPASSIGNEE),
+        ]

snuba/datasets/cdc/groupedmessage.py

@@ -1,7 +1,13 @@
+from typing import Sequence
+
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.entity import Entity
 
 
 class GroupedMessageDataset(Dataset):
-    def __init__(self) -> None:
-        super().__init__(default_entity=EntityKey.GROUPEDMESSAGE)
+    def get_all_entities(self) -> Sequence[Entity]:
+        return [
+            get_entity(EntityKey.GROUPEDMESSAGE),
+        ]

snuba/datasets/configuration/dataset_builder.py

@@ -12,7 +12,6 @@ def build_dataset_from_config(config_file_path: str) -> PluggableDataset:
     config = load_configuration_data(config_file_path, DATASET_VALIDATION_SCHEMAS)
     return PluggableDataset(
         name=config["name"],
-        is_experimental=bool(config["is_experimental"]),
-        default_entity=EntityKey(config["entities"]["default"]),
         all_entities=[EntityKey(key) for key in config["entities"]["all"]],
+        is_experimental=bool(config["is_experimental"]),
     )

snuba/datasets/configuration/generic_metrics/dataset.yaml

@@ -5,5 +5,4 @@ name: generic_metrics
 is_experimental: false
 
 entities:
-  default: generic_metrics_sets
   all: [generic_metrics_sets, generic_metrics_distributions]

snuba/datasets/configuration/json_schema.py

@@ -228,9 +228,13 @@ def make_column_schema(
         "is_experimental": {"type": "boolean"},
         "entities": {
             "type": "object",
-            "properties": {"default": TYPE_STRING, "all": TYPE_STRING_ARRAY},
+            "properties": {"all": TYPE_STRING_ARRAY},
+            "required": ["all"],
         },
     },
+    "required": [
+        "entities",
+    ],
 }
 
 V1_ENTITY_SUBSCIPTION_SCHEMA = {

snuba/datasets/dataset.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+from abc import abstractmethod
 from typing import Sequence
 
-from snuba.datasets.entities.entity_key import EntityKey
 from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.entity import Entity
 from snuba.datasets.plans.query_plan import QueryRunner
@@ -40,9 +40,6 @@ class Dataset:
     manipulate the lower layer objects.
     """
 
-    def __init__(self, *, default_entity: EntityKey) -> None:
-        self.__default_entity = default_entity
-
     def is_experimental(self) -> bool:
         """Marks the dataset as experimental. Healthchecks failing on this
         dataset:
@@ -52,15 +49,9 @@ def is_experimental(self) -> bool:
         """
         return False
 
-    # TODO: Remove once entity selection moves to Sentry
-    def select_entity(self, query: Query) -> EntityKey:
-        return self.__default_entity
-
-    def get_default_entity(self) -> Entity:
-        return get_entity(self.__default_entity)
-
+    @abstractmethod
     def get_all_entities(self) -> Sequence[Entity]:
-        return [self.get_default_entity()]
+        return []
 
     def get_query_pipeline_builder(self) -> DatasetQueryPipelineBuilder:
         return DatasetQueryPipelineBuilder()

snuba/datasets/discover.py

@@ -1,10 +1,13 @@
 import logging
+from typing import Sequence
 
 from snuba import environment
 from snuba.clickhouse.columns import ColumnSet
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.discover import EVENTS_COLUMNS, TRANSACTIONS_COLUMNS
 from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.entity import Entity
 from snuba.query.conditions import (
     BINARY_OPERATORS,
     ConditionFunctions,
@@ -26,9 +29,6 @@
 
 
 class DiscoverDataset(Dataset):
-    def __init__(self) -> None:
-        super().__init__(default_entity=EntityKey.DISCOVER)
-
     # XXX: This is temporary code that will eventually need to be ported to Sentry
     # since SnQL will require an entity to always be specified by the user.
     def select_entity(self, query: Query) -> EntityKey:
@@ -40,6 +40,11 @@ def select_entity(self, query: Query) -> EntityKey:
 
         return selected_entity
 
+    def get_all_entities(self) -> Sequence[Entity]:
+        return [
+            get_entity(EntityKey.DISCOVER),
+        ]
+
 
 metrics = MetricsWrapper(environment.metrics, "api.discover")
 logger = logging.getLogger(__name__)

snuba/datasets/events.py

@@ -1,7 +1,13 @@
+from typing import Sequence
+
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.entity import Entity
 
 
 class EventsDataset(Dataset):
-    def __init__(self) -> None:
-        super().__init__(default_entity=EntityKey.EVENTS)
+    def get_all_entities(self) -> Sequence[Entity]:
+        return [
+            get_entity(EntityKey.EVENTS),
+        ]

snuba/datasets/functions.py

@@ -1,11 +1,17 @@
+from typing import Sequence
+
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.entity import Entity
 
 
 class FunctionsDataset(Dataset):
-    def __init__(self) -> None:
-        super().__init__(default_entity=EntityKey.FUNCTIONS)
-
     @classmethod
     def is_experimental(cls) -> bool:
         return True
+
+    def get_all_entities(self) -> Sequence[Entity]:
+        return [
+            get_entity(EntityKey.FUNCTIONS),
+        ]