Address comments

MetaphorData · Aug 30, 2024 · 18aac85 · 18aac85
1 parent 3e2d1f9
commit 18aac85
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/metaphor/dbt/cloud/utils.py b/metaphor/dbt/cloud/utils.py
@@ -3,16 +3,19 @@
 )
 from metaphor.models.metadata_change_event import DataPlatform
 
+ADAPTER_TYPE_MAPPING = {
+    "DATABRICKS": DataPlatform.UNITY_CATALOG,
+    "POSTGRES": DataPlatform.POSTGRESQL,
+}
+
 
 def parse_environment(environment: GetEnvironmentAdapterTypeEnvironment):
     adapter_type = (
         environment.adapter_type or "unknown"
     )  # It's possible for the environment to not have an adapter type!
     adapter_type = adapter_type.upper()
-    if adapter_type == "DATABRICKS":
-        platform = DataPlatform.UNITY_CATALOG
-    elif adapter_type == "POSTGRES":
-        platform = DataPlatform.POSTGRESQL
+    if adapter_type in ADAPTER_TYPE_MAPPING:
+        platform = ADAPTER_TYPE_MAPPING[adapter_type]
     else:
         assert (
             adapter_type in DataPlatform.__members__