refactor(ingest): add ALL_ENV_TYPES constant (#5866)

datahub-project · Sep 8, 2022 · da46dcd · da46dcd
1 parent 6063484
commit da46dcd
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 36 deletions.
diff --git a/metadata-ingestion/src/datahub/configuration/source_common.py b/metadata-ingestion/src/datahub/configuration/source_common.py
@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, Set
 
 from pydantic import validator
 from pydantic.fields import Field
@@ -8,6 +8,11 @@
 
 DEFAULT_ENV = FabricTypeClass.PROD
 
+# Get all the constants from the FabricTypeClass. It's not an enum, so this is a bit hacky but works.
+ALL_ENV_TYPES: Set[str] = set(
+    [value for name, value in vars(FabricTypeClass).items() if not name.startswith("_")]
+)
+
 
 class PlatformSourceConfigBase(ConfigModel):
     """
@@ -30,20 +35,14 @@ class EnvBasedSourceConfigBase(ConfigModel):
     """
 
     env: str = Field(
-        default=FabricTypeClass.PROD,
+        default=DEFAULT_ENV,
         description="The environment that all assets produced by this connector belong to",
     )
 
     @validator("env")
     def env_must_be_one_of(cls, v: str) -> str:
-        # Get all the constants from the FabricTypeClass. It's not an enum, so this is a bit hacky but works
-        allowed_envs = [
-            value
-            for name, value in vars(FabricTypeClass).items()
-            if not name.startswith("_")
-        ]
-        if (v.upper()) not in allowed_envs:
-            raise ConfigurationError(f"env must be one of {allowed_envs}, found {v}")
+        if v.upper() not in ALL_ENV_TYPES:
+            raise ConfigurationError(f"env must be one of {ALL_ENV_TYPES}, found {v}")
         return v.upper()
 
 

diff --git a/metadata-ingestion/src/datahub/utilities/urns/data_flow_urn.py b/metadata-ingestion/src/datahub/utilities/urns/data_flow_urn.py
@@ -1,6 +1,6 @@
-from typing import List, Set
+from typing import List
 
-from datahub.metadata.schema_classes import FabricTypeClass
+from datahub.configuration.source_common import ALL_ENV_TYPES
 from datahub.utilities.urns.error import InvalidUrnError
 from datahub.utilities.urns.urn import Urn
 
@@ -12,13 +12,6 @@ class DataFlowUrn(Urn):
     """
 
     ENTITY_TYPE: str = "dataFlow"
-    VALID_FABRIC_SET: Set[str] = set(
-        [
-            str(getattr(FabricTypeClass, attr)).upper()
-            for attr in dir(FabricTypeClass)
-            if not callable(getattr(FabricTypeClass, attr)) and not attr.startswith("_")
-        ]
-    )
 
     def __init__(
         self, entity_type: str, entity_id: List[str], domain: str = Urn.LI_DOMAIN
@@ -81,7 +74,5 @@ def _validate_entity_id(entity_id: List[str]) -> None:
             )
 
         env = entity_id[2].upper()
-        if env not in DataFlowUrn.VALID_FABRIC_SET:
-            raise InvalidUrnError(
-                f"Invalid env:{env}. Allowed evn are {DataFlowUrn.VALID_FABRIC_SET}"
-            )
+        if env not in ALL_ENV_TYPES:
+            raise InvalidUrnError(f"Invalid env:{env}. Allowed evn are {ALL_ENV_TYPES}")
diff --git a/metadata-ingestion/src/datahub/utilities/urns/dataset_urn.py b/metadata-ingestion/src/datahub/utilities/urns/dataset_urn.py
@@ -1,6 +1,6 @@
-from typing import List, Optional, Set
+from typing import List, Optional
 
-from datahub.metadata.schema_classes import FabricTypeClass
+from datahub.configuration.source_common import ALL_ENV_TYPES
 from datahub.utilities.urns.data_platform_urn import DataPlatformUrn
 from datahub.utilities.urns.error import InvalidUrnError
 from datahub.utilities.urns.urn import Urn
@@ -13,13 +13,6 @@ class DatasetUrn(Urn):
     """
 
     ENTITY_TYPE: str = "dataset"
-    VALID_FABRIC_SET: Set[str] = set(
-        [
-            str(getattr(FabricTypeClass, attr)).upper()
-            for attr in dir(FabricTypeClass)
-            if not callable(getattr(FabricTypeClass, attr)) and not attr.startswith("_")
-        ]
-    )
 
     def __init__(self, entity_type: str, entity_id: List[str], domain: str = "li"):
         super().__init__(entity_type, entity_id, domain)
@@ -95,10 +88,8 @@ def _validate_entity_id(entity_id: List[str]) -> None:
 
         DataPlatformUrn.validate(platform_urn_str)
         env = entity_id[2].upper()
-        if env not in DatasetUrn.VALID_FABRIC_SET:
-            raise InvalidUrnError(
-                f"Invalid env:{env}. Allowed evn are {DatasetUrn.VALID_FABRIC_SET}"
-            )
+        if env not in ALL_ENV_TYPES:
+            raise InvalidUrnError(f"Invalid env:{env}. Allowed evn are {ALL_ENV_TYPES}")
 
     """A helper function to extract simple . path notation from the v2 field path"""