From ed267642b0e9ce5f6b9e41189535303fa6dc3f87 Mon Sep 17 00:00:00 2001
From: Mike Gouline <1960272+gouline@users.noreply.github.com>
Date: Fri, 1 Nov 2024 22:53:54 +1100
Subject: [PATCH] support relationships to sources

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

diff --git a/dbtmetabase/manifest.py b/dbtmetabase/manifest.py
index 014ef03..61a3dfd 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,14 @@ class Group(str, Enum):
     nodes = "nodes"
     sources = "sources"
 
+    @staticmethod
+    def from_unique_id(unique_id: str) -> Optional[Group]:
+        if prefix := unique_id.split(".")[0] == "source":
+            return Group.sources
+        elif prefix == "model":
+            return Group.nodes
+        return None
+
 
 @dc.dataclass
 class Column: