Update method names after rebasing

gretelai · May 16, 2023 · 5aeb8a5 · 5aeb8a5
1 parent 41c5852
commit 5aeb8a5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/gretel_trainer/relational/core.py b/src/gretel_trainer/relational/core.py
@@ -137,14 +137,14 @@ def add_table(
 
     def _add_rel_json_and_tables(self, table: str, rj_ingest: IngestResponseT) -> None:
         rel_json, commands = rj_ingest
-        add_tables, add_foreign_keys = commands
+        tables, foreign_keys = commands
 
         self.relational_jsons[table] = rel_json
 
-        for add_table in add_tables:
-            self._add_single_table(**add_table)
-        for add_foreign_key in add_foreign_keys:
-            self.add_foreign_key(**add_foreign_key)
+        for tbl in tables:
+            self._add_single_table(**tbl)
+        for foreign_key in foreign_keys:
+            self.add_foreign_key_constraint(**foreign_key)
 
     def _add_single_table(
         self,
@@ -193,7 +193,7 @@ def set_primary_key(
 
             self._add_rel_json_and_tables(table, new_rj_ingest)
             for fk in original_fks:
-                self.add_foreign_key(
+                self.add_foreign_key_constraint(
                     table=fk.table_name,
                     constrained_columns=fk.columns,
                     referred_table=fk.parent_table_name,
@@ -386,7 +386,7 @@ def update_table_data(self, table: str, data: pd.DataFrame) -> None:
                 )
                 parent_table_name = table
             for fk in original_fks:
-                self.add_foreign_key(
+                self.add_foreign_key_constraint(
                     table=fk.table_name,
                     constrained_columns=fk.columns,
                     referred_table=parent_table_name,
@@ -409,7 +409,7 @@ def update_table_data(self, table: str, data: pd.DataFrame) -> None:
                 self.graph.remove_node(table)
                 self._add_rel_json_and_tables(table, new_rj_ingest)
                 for fk in original_foreign_keys:
-                    self.add_foreign_key(
+                    self.add_foreign_key_constraint(
                         table=fk.table_name,
                         constrained_columns=fk.columns,
                         referred_table=new_rj_ingest[0].root_table_name,

diff --git a/tests/relational/test_relational_data_with_json.py b/tests/relational/test_relational_data_with_json.py
@@ -234,7 +234,9 @@ def test_foreign_keys(documents):
     ]
 
     # Removing a foreign key from the source-with-JSON table updates the root invented table
-    documents.remove_foreign_key(table="purchases", constrained_columns=["user_id"])
+    documents.remove_foreign_key_constraint(
+        table="purchases", constrained_columns=["user_id"]
+    )
     assert documents.get_foreign_keys("purchases") == []
     assert documents.get_foreign_keys("purchases-sfx") == []
 
@@ -361,13 +363,13 @@ def test_update_data_existing_flat_to_json(documents):
     rel_data.add_table(
         name="payments", primary_key="id", data=documents.get_table_data("payments")
     )
-    rel_data.add_foreign_key(
+    rel_data.add_foreign_key_constraint(
         table="purchases",
         constrained_columns=["user_id"],
         referred_table="users",
         referred_columns=["id"],
     )
-    rel_data.add_foreign_key(
+    rel_data.add_foreign_key_constraint(
         table="payments",
         constrained_columns=["purchase_id"],
         referred_table="purchases",