From ee2ab6cd0e8a54117880fa64f671fcb99d2b9875 Mon Sep 17 00:00:00 2001 From: MrCurtis <4184070+MrCurtis@users.noreply.github.com> Date: Thu, 9 Nov 2023 19:41:11 +0000 Subject: [PATCH] Ensure rows added for only selected foreign-keys. Previously, when building the graph, the code would only add rows realted through a reverse foreign-key relation only if they were from a table included in the `only_tables` argument. However, rows related by forward foreign-keys would be added regardless. This was a mistake, and now the forward case is similar to the reverse case. --- graph.py | 8 ++++---- pyproject.toml | 2 +- test_graph.py | 18 +++++++++++++++++- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/graph.py b/graph.py index a839984..7ec57c2 100644 --- a/graph.py +++ b/graph.py @@ -84,7 +84,7 @@ def get_graph(engine, table, primary_key, only_tables=None) -> Graph: data=_get_data(row), ) graph.add_node(row_node) - _add_related_rows_to_graph(row, row_node, graph) + _add_related_rows_to_graph(row, row_node, graph, only_tables) return graph @@ -107,7 +107,7 @@ def _get_row(session, table, primary_key): ) -def _add_related_rows_to_graph(row, row_node, graph): +def _add_related_rows_to_graph(row, row_node, graph, only_tables): related = [] relationships = row.__mapper__.relationships for relationship in relationships: @@ -127,7 +127,7 @@ def _add_related_rows_to_graph(row, row_node, graph): # This path for foreign keys. related_row = related_rows # Ignore null foreign-keys. - if related_row is not None: + if related_row is not None and (only_tables is None or related_row.__table__.name in only_tables): related_node = Node( table=_get_table_name_from_row(related_row), primary_key=_get_primary_key_from_row(related_row), @@ -141,7 +141,7 @@ def _add_related_rows_to_graph(row, row_node, graph): for _, related_node in related: graph.add_edge(row_node, related_node) for unvisited_row, unvisited_node in unvisited: - _add_related_rows_to_graph(unvisited_row, unvisited_node, graph) + _add_related_rows_to_graph(unvisited_row, unvisited_node, graph, only_tables) def _get_table_name_from_row(row): diff --git a/pyproject.toml b/pyproject.toml index c530e5d..fe80204 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [project] name = "fk-graph" -version = "0.0.5" +version = "0.0.6" authors = [ { name="Andrew Curtis", email="fk.graph@fastmail.com" }, { name="John C Thomas" }, diff --git a/test_graph.py b/test_graph.py index 99b6c2d..63224ca 100644 --- a/test_graph.py +++ b/test_graph.py @@ -135,7 +135,7 @@ def test_raises_primary_key_does_not_exist_when_no_key_in_table(self): ): get_graph(self.engine, "table_a", "9876") - def test_can_restrict_to_selected_tables(self): + def test_can_restrict_to_selected_tables__reverse_foreign_key_case(self): self._create_three_entries_with_linear_foreign_key_relations(self.engine) graph = get_graph(self.engine, "table_a", "1", only_tables=["table_a", "table_b"]) @@ -151,6 +151,22 @@ def test_can_restrict_to_selected_tables(self): any([n.table == "table_c" for n in graph.nodes]) ) + def test_can_restrict_to_selected_tables__forward_foreign_key_case(self): + self._create_three_entries_with_linear_foreign_key_relations(self.engine) + + graph = get_graph(self.engine, "table_c", "1", only_tables=["table_c", "table_b"]) + + with self.subTest("includes selected tables"): + self.assertTrue( + any([n.table == "table_c" for n in graph.nodes]) + and + any([n.table == "table_b" for n in graph.nodes]) + ) + with self.subTest("excludes non-selected tables"): + self.assertFalse( + any([n.table == "table_a" for n in graph.nodes]) + ) + def test_can_create_graph_when_some_rows_have_null_foreign_keys(self): self._create_entries_with_null_foreign_key(self.engine)