Store entities keys in dataset class

enochtangg · enochtangg · commit 9c20561f20bf · 2022-09-15T11:30:11.000-04:00
diff --git a/snuba/datasets/cdc/groupassignee.py b/snuba/datasets/cdc/groupassignee.py
@@ -1,6 +1,5 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class GroupAssigneeDataset(Dataset):
@@ -15,6 +14,6 @@ class GroupAssigneeDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.GROUPASSIGNEE),
+                EntityKey.GROUPASSIGNEE,
             ]
         )
diff --git a/snuba/datasets/cdc/groupedmessage.py b/snuba/datasets/cdc/groupedmessage.py
@@ -1,12 +1,11 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class GroupedMessageDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.GROUPEDMESSAGE),
+                EntityKey.GROUPEDMESSAGE,
             ]
         )
diff --git a/snuba/datasets/dataset.py b/snuba/datasets/dataset.py
@@ -2,6 +2,7 @@
 
 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
@@ -39,7 +40,7 @@ class Dataset:
     manipulate the lower layer objects.
     """
 
-    def __init__(self, *, all_entities: Sequence[Entity]) -> None:
+    def __init__(self, *, all_entities: Sequence[EntityKey]) -> None:
         self.__all_entities = all_entities
 
     def is_experimental(self) -> bool:
@@ -52,7 +53,7 @@ def is_experimental(self) -> bool:
         return False
 
     def get_all_entities(self) -> Sequence[Entity]:
-        return self.__all_entities
+        return [get_entity(entity_key) for entity_key in self.__all_entities]
 
     def get_query_pipeline_builder(self) -> DatasetQueryPipelineBuilder:
         return DatasetQueryPipelineBuilder()
diff --git a/snuba/datasets/discover.py b/snuba/datasets/discover.py
@@ -5,7 +5,6 @@
 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.query.conditions import (
     BINARY_OPERATORS,
     ConditionFunctions,
@@ -30,9 +29,9 @@ class DiscoverDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.DISCOVER),
-                get_entity(EntityKey.DISCOVER_EVENTS),
-                get_entity(EntityKey.DISCOVER_TRANSACTIONS),
+                EntityKey.DISCOVER,
+                EntityKey.DISCOVER_EVENTS,
+                EntityKey.DISCOVER_TRANSACTIONS,
             ]
         )
 
diff --git a/snuba/datasets/events.py b/snuba/datasets/events.py
@@ -1,12 +1,11 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class EventsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.EVENTS),
+                EntityKey.EVENTS,
             ]
         )
diff --git a/snuba/datasets/functions.py b/snuba/datasets/functions.py
@@ -1,13 +1,12 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class FunctionsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.FUNCTIONS),
+                EntityKey.FUNCTIONS,
             ]
         )
 
diff --git a/snuba/datasets/generic_metrics.py b/snuba/datasets/generic_metrics.py
@@ -1,6 +1,5 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 DEFAULT_GRANULARITY = 60
 
@@ -9,7 +8,7 @@ class GenericMetricsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.GENERIC_METRICS_SETS),
-                get_entity(EntityKey.GENERIC_METRICS_DISTRIBUTIONS),
+                EntityKey.GENERIC_METRICS_SETS,
+                EntityKey.GENERIC_METRICS_DISTRIBUTIONS,
             ]
         )
diff --git a/snuba/datasets/metrics.py b/snuba/datasets/metrics.py
@@ -1,15 +1,14 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class MetricsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.METRICS_COUNTERS),
-                get_entity(EntityKey.METRICS_DISTRIBUTIONS),
-                get_entity(EntityKey.METRICS_SETS),
-                get_entity(EntityKey.ORG_METRICS_COUNTERS),
+                EntityKey.METRICS_COUNTERS,
+                EntityKey.METRICS_DISTRIBUTIONS,
+                EntityKey.METRICS_SETS,
+                EntityKey.ORG_METRICS_COUNTERS,
             ]
         )
diff --git a/snuba/datasets/outcomes.py b/snuba/datasets/outcomes.py
@@ -1,6 +1,5 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class OutcomesDataset(Dataset):
@@ -11,6 +10,6 @@ class OutcomesDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.OUTCOMES),
+                EntityKey.OUTCOMES,
             ]
         )
diff --git a/snuba/datasets/outcomes_raw.py b/snuba/datasets/outcomes_raw.py
@@ -1,6 +1,5 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class OutcomesRawDataset(Dataset):
@@ -11,6 +10,6 @@ class OutcomesRawDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.OUTCOMES_RAW),
+                EntityKey.OUTCOMES_RAW,
             ]
         )
diff --git a/snuba/datasets/pluggable_dataset.py b/snuba/datasets/pluggable_dataset.py
@@ -4,7 +4,6 @@
 
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class PluggableDataset(Dataset):
@@ -22,9 +21,7 @@ def __init__(
         all_entities: list[EntityKey],
         is_experimental: bool | None,
     ) -> None:
-        super().__init__(
-            all_entities=[get_entity(entity_key) for entity_key in all_entities]
-        )
+        super().__init__(all_entities=all_entities)
         self.name = name
         self.__is_experimental = is_experimental or False
 
diff --git a/snuba/datasets/profiles.py b/snuba/datasets/profiles.py
@@ -1,12 +1,11 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class ProfilesDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.PROFILES),
+                EntityKey.PROFILES,
             ]
         )
diff --git a/snuba/datasets/replays.py b/snuba/datasets/replays.py
@@ -1,13 +1,12 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class ReplaysDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.REPLAYS),
+                EntityKey.REPLAYS,
             ]
         )
 
diff --git a/snuba/datasets/sessions.py b/snuba/datasets/sessions.py
@@ -1,12 +1,11 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class SessionsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.SESSIONS),
+                EntityKey.SESSIONS,
             ]
         )
diff --git a/snuba/datasets/transactions.py b/snuba/datasets/transactions.py
@@ -1,12 +1,11 @@
 from snuba.datasets.dataset import Dataset
 from snuba.datasets.entities.entity_key import EntityKey
-from snuba.datasets.entities.factory import get_entity
 
 
 class TransactionsDataset(Dataset):
     def __init__(self) -> None:
         super().__init__(
             all_entities=[
-                get_entity(EntityKey.TRANSACTIONS),
+                EntityKey.TRANSACTIONS,
             ]
         )

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,5 @@`
`1`	`1`	`from snuba.datasets.dataset import Dataset`
`2`	`2`	`from snuba.datasets.entities.entity_key import EntityKey`
`3`		`-from snuba.datasets.entities.factory import get_entity`
`4`	`3`
`5`	`4`
`6`	`5`	`class GroupAssigneeDataset(Dataset):`
`@@ -15,6 +14,6 @@ class GroupAssigneeDataset(Dataset):`
`15`	`14`	`def __init__(self) -> None:`
`16`	`15`	`super().__init__(`
`17`	`16`	`all_entities=[`
`18`		`- get_entity(EntityKey.GROUPASSIGNEE),`
	`17`	`+ EntityKey.GROUPASSIGNEE,`
`19`	`18`	`]`
`20`	`19`	`)`
Original file line number	Diff line number	Diff line change
`@@ -1,12 +1,11 @@`
`1`	`1`	`from snuba.datasets.dataset import Dataset`
`2`	`2`	`from snuba.datasets.entities.entity_key import EntityKey`
`3`		`-from snuba.datasets.entities.factory import get_entity`
`4`	`3`
`5`	`4`
`6`	`5`	`class GroupedMessageDataset(Dataset):`
`7`	`6`	`def __init__(self) -> None:`
`8`	`7`	`super().__init__(`
`9`	`8`	`all_entities=[`
`10`		`- get_entity(EntityKey.GROUPEDMESSAGE),`
	`9`	`+ EntityKey.GROUPEDMESSAGE,`
`11`	`10`	`]`
`12`	`11`	`)`
Original file line number	Diff line number	Diff line change
`@@ -1,13 +1,12 @@`
`1`	`1`	`from snuba.datasets.dataset import Dataset`
`2`	`2`	`from snuba.datasets.entities.entity_key import EntityKey`
`3`		`-from snuba.datasets.entities.factory import get_entity`
`4`	`3`
`5`	`4`
`6`	`5`	`class FunctionsDataset(Dataset):`
`7`	`6`	`def __init__(self) -> None:`
`8`	`7`	`super().__init__(`
`9`	`8`	`all_entities=[`
`10`		`- get_entity(EntityKey.FUNCTIONS),`
	`9`	`+ EntityKey.FUNCTIONS,`
`11`	`10`	`]`
`12`	`11`	`)`
`13`	`12`