Backup cols (#129)

* Remove unused, backwards-incompat columns arg on backup

* Provide explicit columns when writing source archive

* Don't show deprecation warning during restore

Author: mikeknep

src/gretel_trainer/relational/backup.py

@@ -10,7 +10,6 @@
 @dataclass
 class BackupRelationalDataTable:
     primary_key: list[str]
-    columns: list[str]
     invented_table_metadata: Optional[dict[str, Any]] = None
 
 
@@ -53,7 +52,6 @@ def from_relational_data(cls, rel_data: RelationalData) -> BackupRelationalData:
         for table in rel_data.list_all_tables():
             backup_table = BackupRelationalDataTable(
                 primary_key=rel_data.get_primary_key(table),
-                columns=rel_data.get_table_columns(table),
             )
             if (
                 invented_table_metadata := rel_data.get_invented_table_metadata(table)

src/gretel_trainer/relational/multi_table.py

@@ -103,10 +103,11 @@ def __init__(
     ):
         self._strategy = _validate_strategy(strategy)
         if gretel_model is not None:
-            logger.warning(
-                "The `gretel_model` argument is deprecated and will be removed in a future release. "
-                "Going forward you should provide a config to `train_synthetics`."
-            )
+            if backup is None:
+                logger.warning(
+                    "The `gretel_model` argument is deprecated and will be removed in a future release. "
+                    "Going forward you should provide a config to `train_synthetics`."
+                )
             model_name, model_config = self._validate_gretel_model(gretel_model)
             self._gretel_model = model_name
             self._model_config = model_config
@@ -945,7 +946,11 @@ def _upload_sources_to_project(self) -> None:
                 filename = f"source_{table}.csv"
                 out_path = self._working_dir / filename
                 df = self.relational_data.get_table_data(table)
-                df.to_csv(out_path, index=False)
+                df.to_csv(
+                    out_path,
+                    index=False,
+                    columns=self.relational_data.get_table_columns(table),
+                )
                 tar.add(out_path, arcname=filename)
         self._artifact_collection.upload_source_archive(
             self._project, str(archive_path)

tests/relational/test_backup.py

@@ -19,11 +19,9 @@ def test_backup_relational_data(trips):
         tables={
             "vehicle_types": BackupRelationalDataTable(
                 primary_key=["id"],
-                columns=["id", "name"],
             ),
             "trips": BackupRelationalDataTable(
                 primary_key=["id"],
-                columns=["id", "purpose", "vehicle_type_id"],
             ),
         },
         foreign_keys=[
@@ -43,19 +41,9 @@ def test_backup_relational_data(trips):
 def test_backup_relational_data_with_json(documents):
     expected = BackupRelationalData(
         tables={
-            "users": BackupRelationalDataTable(
-                primary_key=["id"], columns=["id", "name"]
-            ),
+            "users": BackupRelationalDataTable(primary_key=["id"]),
             "purchases-sfx": BackupRelationalDataTable(
                 primary_key=["id", "~PRIMARY_KEY_ID~"],
-                columns=[
-                    "~PRIMARY_KEY_ID~",
-                    "id",
-                    "user_id",
-                    "data>item",
-                    "data>cost",
-                    "data>details>color",
-                ],
                 invented_table_metadata={
                     "invented_root_table_name": "purchases-sfx",
                     "original_table_name": "purchases",
@@ -64,16 +52,13 @@ def test_backup_relational_data_with_json(documents):
             ),
             "purchases-data-years-sfx": BackupRelationalDataTable(
                 primary_key=["~PRIMARY_KEY_ID~"],
-                columns=["~PRIMARY_KEY_ID~", "purchases~id", "content", "array~order"],
                 invented_table_metadata={
                     "invented_root_table_name": "purchases-sfx",
                     "original_table_name": "purchases",
                     "empty": False,
                 },
             ),
-            "payments": BackupRelationalDataTable(
-                primary_key=["id"], columns=["id", "purchase_id", "amount"]
-            ),
+            "payments": BackupRelationalDataTable(primary_key=["id"]),
         },
         foreign_keys=[
             BackupForeignKey(
@@ -116,11 +101,9 @@ def test_backup():
         tables={
             "customer": BackupRelationalDataTable(
                 primary_key=["id"],
-                columns=["a", "b", "c"],
             ),
             "address": BackupRelationalDataTable(
                 primary_key=[],
-                columns=["x", "y", "z"],
             ),
         },
         foreign_keys=[