Support for JSON (#107)

requirements.txt

@@ -10,3 +10,4 @@ requests~=2.25
 scikit-learn~=1.0
 smart-open[s3]~=5.2
 sqlalchemy~=1.4
+unflatten==0.1.1

src/gretel_trainer/relational/__init__.py

@@ -8,4 +8,5 @@
     sqlite_conn,
 )
 from gretel_trainer.relational.core import RelationalData
+from gretel_trainer.relational.log import set_log_level
 from gretel_trainer.relational.multi_table import MultiTable

src/gretel_trainer/relational/backup.py

@@ -10,6 +10,7 @@
 @dataclass
 class BackupRelationalDataTable:
     primary_key: List[str]
+    invented_table_metadata: Optional[dict[str, str]] = None
 
 
 @dataclass
@@ -29,26 +30,52 @@ def from_fk(cls, fk: ForeignKey) -> BackupForeignKey:
         )
 
 
+@dataclass
+class BackupRelationalJson:
+    original_table_name: str
+    original_primary_key: list[str]
+    original_columns: list[str]
+    table_name_mappings: dict[str, str]
+    invented_table_names: list[str]
+
+
 @dataclass
 class BackupRelationalData:
     tables: Dict[str, BackupRelationalDataTable]
     foreign_keys: List[BackupForeignKey]
+    relational_jsons: Dict[str, BackupRelationalJson]
 
     @classmethod
     def from_relational_data(cls, rel_data: RelationalData) -> BackupRelationalData:
         tables = {}
         foreign_keys = []
+        relational_jsons = {}
         for table in rel_data.list_all_tables():
-            tables[table] = BackupRelationalDataTable(
+            backup_table = BackupRelationalDataTable(
                 primary_key=rel_data.get_primary_key(table),
             )
+            if (
+                invented_table_metadata := rel_data.get_invented_table_metadata(table)
+            ) is not None:
+                backup_table.invented_table_metadata = asdict(invented_table_metadata)
+            tables[table] = backup_table
             foreign_keys.extend(
                 [
                     BackupForeignKey.from_fk(key)
                     for key in rel_data.get_foreign_keys(table)
                 ]
             )
-        return BackupRelationalData(tables=tables, foreign_keys=foreign_keys)
+        for key, rel_json in rel_data.relational_jsons.items():
+            relational_jsons[key] = BackupRelationalJson(
+                original_table_name=rel_json.original_table_name,
+                original_primary_key=rel_json.original_primary_key,
+                original_columns=rel_json.original_columns,
+                table_name_mappings=rel_json.table_name_mappings,
+                invented_table_names=rel_json.table_names,
+            )
+        return BackupRelationalData(
+            tables=tables, foreign_keys=foreign_keys, relational_jsons=relational_jsons
+        )
 
 
 @dataclass
@@ -114,6 +141,10 @@ def from_dict(cls, b: Dict[str, Any]):
                 )
                 for fk in relational_data.get("foreign_keys", [])
             ],
+            relational_jsons={
+                k: BackupRelationalJson(**v)
+                for k, v in relational_data.get("relational_jsons", {}).items()
+            },
         )
 
         backup = Backup(