feat: allow self reference (#121)

* feat: allow self reference * chore: review comments * fix: explicit test for node in list before removal Co-authored-by: Mike Gouline <[email protected]>
gouline · Jul 9, 2022 · de5574c · de5574c
1 parent b03d1d3
commit de5574c
Showing 1 changed file with 26 additions and 6 deletions.
diff --git a/dbtmetabase/parsers/dbt_manifest.py b/dbtmetabase/parsers/dbt_manifest.py
@@ -185,7 +185,9 @@ def _read_model(
 
         metabase_column: List[MetabaseColumn] = []
 
-        children = manifest["child_map"][model["unique_id"]]
+        unique_id = model["unique_id"]
+
+        children = manifest["child_map"][unique_id]
         relationship_tests = {}
 
         for child_id in children:
@@ -201,12 +203,31 @@ def _read_model(
                 # would return the ref() written in the test, but if the model has an alias, that's not enough.
                 # It is better to use child['depends_on']['nodes'] and exclude the current model
 
-                depends_on_ids = set(child["depends_on"][model_type])
-                depends_on_ids.discard(model["unique_id"])
-                if not depends_on_ids:
+                # From experience, nodes contains at most two tables: the referenced model and the current model.
+                # Note, sometimes only the referenced model is returned.
+                depends_on_nodes = list(child["depends_on"][model_type])
+                if len(depends_on_nodes) > 2:
+                    logger().warning(
+                        "Expected at most two nodes, got %d {} nodes, skipping %s {}",
+                        len(depends_on_nodes),
+                        unique_id,
+                    )
+                    continue
+
+                # Remove the current model from the list. Note, remove() only removes the first occurrence. This ensures
+                # the logic also works for self referencing models.
+                if len(depends_on_nodes) == 2 and unique_id in depends_on_nodes:
+                    depends_on_nodes.remove(unique_id)
+
+                if len(depends_on_nodes) != 1:
+                    logger().warning(
+                        "Expected single node after filtering, got %d nodes, skipping %s",
+                        len(depends_on_nodes),
+                        unique_id,
+                    )
                     continue
 
-                depends_on_id = depends_on_ids.pop()
+                depends_on_id = depends_on_nodes[0]
 
                 foreign_key_model = manifest[model_type].get(depends_on_id, {})
                 fk_target_table_alias = foreign_key_model.get(
@@ -252,7 +273,6 @@ def _read_model(
                     description += "\n\n"
                 description += f"Tags: {tags}"
 
-        unique_id = model["unique_id"]
         if docs_url:
             full_path = f"{docs_url}/#!/model/{unique_id}"
             if description != "":