Align relational outputs

GitOrigin-RevId: 6fa84a0d13575a1c067f56c81b3a6abd9d005cd3
gretelai · Jan 22, 2024 · cc4e406 · cc4e406
1 parent c832c4f
commit cc4e406
Show file tree

Hide file tree

Showing 10 changed files with 157 additions and 965 deletions.
diff --git a/src/gretel_trainer/relational/artifacts.py b/src/gretel_trainer/relational/artifacts.py
diff --git a/src/gretel_trainer/relational/backup.py b/src/gretel_trainer/relational/backup.py
@@ -3,7 +3,6 @@
 from dataclasses import asdict, dataclass
 from typing import Any, Optional, Union
 
-from gretel_trainer.relational.artifacts import ArtifactCollection
 from gretel_trainer.relational.core import ForeignKey, RelationalData, Scope
 from gretel_trainer.relational.json import InventedTableMetadata, ProducerMetadata
 
@@ -104,7 +103,7 @@ class Backup:
     project_name: str
     strategy: str
     refresh_interval: int
-    artifact_collection: ArtifactCollection
+    source_archive: Optional[str]
     relational_data: BackupRelationalData
     classify: Optional[BackupClassify] = None
     transforms_train: Optional[BackupTransformsTrain] = None
@@ -134,11 +133,16 @@ def from_dict(cls, b: dict[str, Any]):
             ],
         )
 
+        # source_archive previously was stored under artifact_collection
+        source_archive = b.get(
+            "source_archive", b.get("artifact_collection", {}).get("source_archive")
+        )
+
         backup = Backup(
             project_name=b["project_name"],
             strategy=b["strategy"],
             refresh_interval=b["refresh_interval"],
-            artifact_collection=ArtifactCollection(**b["artifact_collection"]),
+            source_archive=source_archive,
             relational_data=brd,
         )