From e900f513d7b5b74c4fbc9a14ab3e05e1ad73bd29 Mon Sep 17 00:00:00 2001
From: Mike Gouline <1960272+gouline@users.noreply.github.com>
Date: Fri, 1 Nov 2024 23:09:54 +1100
Subject: [PATCH] support relationships to sources (#289)

---
 dbtmetabase/manifest.py | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/dbtmetabase/manifest.py b/dbtmetabase/manifest.py
index 014ef03..f490849 100644
--- a/dbtmetabase/manifest.py
+++ b/dbtmetabase/manifest.py
@@ -230,19 +230,22 @@ def _read_relationships(
                     continue
 
                 depends_on_id = depends_on_nodes[0]
+                depends_on_group = Group.from_unique_id(depends_on_id)
+                if not depends_on_group:
+                    _logger.debug("Unknown group dependency '%s'", depends_on_id)
+                    continue
 
-                fk_target_model = manifest[group].get(depends_on_id, {})
-                fk_target_table = fk_target_model.get(
-                    "alias",
-                    fk_target_model.get("identifier", fk_target_model.get("name")),
+                fk_target_model = manifest[depends_on_group].get(depends_on_id, {})
+                fk_target_table = (
+                    fk_target_model.get("alias")
+                    or fk_target_model.get("identifier")
+                    or fk_target_model.get("name")
                 )
                 if not fk_target_table:
                     _logger.debug("Cannot resolve dependency for '%s'", depends_on_id)
                     continue
 
-                fk_target_schema = manifest[group][depends_on_id].get(
-                    "schema", DEFAULT_SCHEMA
-                )
+                fk_target_schema = fk_target_model.get("schema", DEFAULT_SCHEMA)
                 fk_target_table = f"{fk_target_schema}.{fk_target_table}"
                 fk_target_field = child["test_metadata"]["kwargs"]["field"].strip('"')
 
@@ -341,6 +344,15 @@ class Group(str, Enum):
     nodes = "nodes"
     sources = "sources"
 
+    @staticmethod
+    def from_unique_id(unique_id: str) -> Optional[Group]:
+        prefix = unique_id.split(".")[0]
+        if prefix == "source":
+            return Group.sources
+        elif prefix == "model":
+            return Group.nodes
+        return None
+
 
 @dc.dataclass
 class Column: