From 18aac8568a137233374d4a8cfb67bdc2f4ee85d8 Mon Sep 17 00:00:00 2001
From: Scott Huang <scott@metaphor.io>
Date: Fri, 30 Aug 2024 21:07:06 +0800
Subject: [PATCH] Address comments

---
 metaphor/dbt/cloud/utils.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/metaphor/dbt/cloud/utils.py b/metaphor/dbt/cloud/utils.py
index 9d58d6c4..e71a9c1a 100644
--- 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__