ingestify 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

ingestify/exceptions.py

@@ -8,3 +8,7 @@ class ConfigurationError(IngestifyError):
 
 class DuplicateFile(IngestifyError):
     pass
+
+
+class SaveError(IngestifyError):
+    pass

ingestify/infra/serialization/__init__.py

@@ -1,50 +1,22 @@
-import json
-from datetime import datetime
-from typing import Type, Any, TypeVar
+from ingestify.domain import DatasetCreated
+from ingestify.domain.models.dataset.events import RevisionAdded
+from ingestify.domain.models.event import DomainEvent
 
-from dataclass_factory import Schema, Factory, NameStyle
-from dataclass_factory.schema_helpers import type_checker
 
-from ingestify.domain import DatasetCreated, Identifier
-from ingestify.domain.models.dataset.events import MetadataUpdated, RevisionAdded
+event_types = {
+    DatasetCreated.event_type: DatasetCreated,
+    RevisionAdded.event_type: RevisionAdded,
+}
 
-isotime_schema = Schema(
-    parser=lambda x: datetime.fromisoformat(x.replace("Z", "+00:00")),  # type: ignore
-    serializer=lambda x: datetime.isoformat(x).replace("+00:00", "Z"),
-)
 
-identifier_schema = Schema(
-    # json.loads(x) for backwards compatibility
-    parser=lambda x: Identifier(x if isinstance(x, dict) else json.loads(x)),
-    serializer=lambda x: dict(x),
-)
+def deserialize(event_dict: dict) -> DomainEvent:
+    event_cls = event_types[event_dict["event_type"]]
+    return event_cls.model_validate(event_dict)
 
-factory = Factory(
-    schemas={
-        datetime: isotime_schema,
-        Identifier: identifier_schema,
-        DatasetCreated: Schema(
-            pre_parse=type_checker(DatasetCreated.event_type, "event_type")
-        ),
-        MetadataUpdated: Schema(
-            pre_parse=type_checker(MetadataUpdated.event_type, "event_type")
-        ),
-        RevisionAdded: Schema(
-            pre_parse=type_checker(RevisionAdded.event_type, "event_type")
-        ),
-        # ClipSelectionContent: Schema(pre_parse=type_checker(ClipSelectionContent.content_type, field="contentType")),
-        # TeamInfoImageContent: Schema(pre_parse=type_checker(TeamInfoImageContent.content_type, field="contentType")),
-        # StaticVideoContent: Schema(pre_parse=type_checker(StaticVideoContent.content_type, field="contentType"))
-    },
-    default_schema=Schema(),
-)
 
-T = TypeVar("T")
+def serialize(event: DomainEvent) -> dict:
+    event_dict = event.model_dump(mode="json")
 
-
-def serialize(data: T, class_: Type[T] = None) -> Any:
-    return factory.dump(data, class_)
-
-
-def unserialize(data: Any, class_: Type[T]) -> T:
-    return factory.load(data, class_)
+    # Make sure event_type is always part of the event_dict. Pydantic might skip it when the type is ClassVar
+    event_dict["event_type"] = event.event_type
+    return event_dict

ingestify/infra/store/dataset/sqlalchemy/repository.py

@@ -1,13 +1,24 @@
+import itertools
 import json
 import uuid
+from collections import defaultdict
 from typing import Optional, Union, List
 
-from sqlalchemy import create_engine, func, text, tuple_
+from sqlalchemy import (
+    create_engine,
+    func,
+    text,
+    tuple_,
+    Table,
+    insert,
+    Transaction,
+    Connection,
+)
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import NoSuchModuleError
 from sqlalchemy.orm import Session, joinedload
 
-from ingestify.domain import File
+from ingestify.domain import File, Revision
 from ingestify.domain.models import (
     Dataset,
     DatasetCollection,
@@ -15,11 +26,22 @@ from ingestify.domain.models import (
     Identifier,
     Selector,
 )
+from ingestify.domain.models.base import BaseModel
 from ingestify.domain.models.dataset.collection_metadata import (
     DatasetCollectionMetadata,
 )
-
-from .mapping import dataset_table, metadata
+from ingestify.domain.models.ingestion.ingestion_job_summary import IngestionJobSummary
+from ingestify.domain.models.task.task_summary import TaskSummary
+from ingestify.exceptions import IngestifyError
+
+from .tables import (
+    metadata,
+    dataset_table,
+    file_table,
+    revision_table,
+    ingestion_job_summary_table,
+    task_summary_table,
+)
 
 
 def parse_value(v):
@@ -113,6 +135,31 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
     def session(self):
         return self.session_provider.get()
 
+    def _upsert(self, connection: Connection, table: Table, entities: list[dict]):
+        dialect = self.session.bind.dialect.name
+        if dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert
+        elif dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert
+        elif dialect == "sqlite":
+            from sqlalchemy.dialects.sqlite import insert
+        else:
+            raise IngestifyError(f"Don't know how to do an upsert in {dialect}")
+
+        stmt = insert(table).values(entities)
+
+        primary_key_columns = [column for column in table.columns if column.primary_key]
+
+        set_ = {
+            name: getattr(stmt.excluded, name)
+            for name, column in table.columns.items()
+            if column not in primary_key_columns
+        }
+
+        stmt = stmt.on_conflict_do_update(index_elements=primary_key_columns, set_=set_)
+
+        connection.execute(stmt)
+
     def _filter_query(
         self,
         query,
@@ -122,11 +169,11 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
         dataset_id: Optional[Union[str, List[str]]] = None,
         selector: Optional[Union[Selector, List[Selector]]] = None,
     ):
-        query = query.filter(Dataset.bucket == bucket)
+        query = query.filter(dataset_table.c.bucket == bucket)
         if dataset_type:
-            query = query.filter(Dataset.dataset_type == dataset_type)
+            query = query.filter(dataset_table.c.dataset_type == dataset_type)
         if provider:
-            query = query.filter(Dataset.provider == provider)
+            query = query.filter(dataset_table.c.provider == provider)
         if dataset_id is not None:
             if isinstance(dataset_id, list):
                 if len(dataset_id) == 0:
@@ -134,9 +181,9 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
                     # return an empty DatasetCollection
                     return DatasetCollection()
 
-                query = query.filter(Dataset.dataset_id.in_(dataset_id))
+                query = query.filter(dataset_table.c.dataset_id.in_(dataset_id))
             else:
-                query = query.filter(Dataset.dataset_id == dataset_id)
+                query = query.filter(dataset_table.c.dataset_id == dataset_id)
 
         dialect = self.session.bind.dialect.name
 
@@ -175,7 +222,7 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
                     else:
                         column = column.as_string()
                 else:
-                    column = func.json_extract(Dataset.identifier, f"$.{k}")
+                    column = func.json_extract(dataset_table.c.identifier, f"$.{k}")
                 columns.append(column)
 
             values = []
@@ -189,6 +236,60 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
             query = query.filter(text(where))
         return query
 
+    def load_datasets(self, dataset_ids: list[str]) -> list[Dataset]:
+        if not dataset_ids:
+            return []
+
+        dataset_rows = list(
+            self.session.query(dataset_table).filter(
+                dataset_table.c.dataset_id.in_(dataset_ids)
+            )
+        )
+        revisions_per_dataset = {}
+        rows = (
+            self.session.query(revision_table)
+            .filter(revision_table.c.dataset_id.in_(dataset_ids))
+            .order_by(revision_table.c.dataset_id)
+        )
+
+        for dataset_id, revisions in itertools.groupby(
+            rows, key=lambda row: row.dataset_id
+        ):
+            revisions_per_dataset[dataset_id] = list(revisions)
+
+        files_per_revision = {}
+        rows = (
+            self.session.query(file_table)
+            .filter(file_table.c.dataset_id.in_(dataset_ids))
+            .order_by(file_table.c.dataset_id, file_table.c.revision_id)
+        )
+
+        for (dataset_id, revision_id), files in itertools.groupby(
+            rows, key=lambda row: (row.dataset_id, row.revision_id)
+        ):
+            files_per_revision[(dataset_id, revision_id)] = list(files)
+
+        datasets = []
+        for dataset_row in dataset_rows:
+            dataset_id = dataset_row.dataset_id
+            revisions = []
+            for revision_row in revisions_per_dataset.get(dataset_id, []):
+                files = [
+                    File.model_validate(file_row)
+                    for file_row in files_per_revision.get(
+                        (dataset_id, revision_row.revision_id), []
+                    )
+                ]
+                revision = Revision.model_validate(
+                    {**revision_row._mapping, "modified_files": files}
+                )
+                revisions.append(revision)
+
+            datasets.append(
+                Dataset.model_validate({**dataset_row._mapping, "revisions": revisions})
+            )
+        return datasets
+
     def get_dataset_collection(
         self,
         bucket: str,
@@ -209,17 +310,19 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
             )
 
         if not metadata_only:
-            dataset_query = apply_query_filter(self.session.query(Dataset))
-            datasets = list(dataset_query)
+            dataset_query = apply_query_filter(
+                self.session.query(dataset_table.c.dataset_id)
+            )
+            dataset_ids = [row.dataset_id for row in dataset_query]
+            datasets = self.load_datasets(dataset_ids)
         else:
             datasets = []
 
         metadata_result_row = apply_query_filter(
             self.session.query(
-                func.min(File.modified_at).label("first_modified_at"),
-                func.max(File.modified_at).label("last_modified_at"),
+                func.max(dataset_table.c.last_modified_at).label("last_modified_at"),
                 func.count().label("row_count"),
-            ).join(Dataset, Dataset.dataset_id == File.dataset_id)
+            )
         ).first()
         dataset_collection_metadata = DatasetCollectionMetadata(*metadata_result_row)
 
@@ -228,12 +331,153 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
     def save(self, bucket: str, dataset: Dataset):
         # Just make sure
         dataset.bucket = bucket
-        self.session.add(dataset)
-        self.session.commit()
+
+        self._save([dataset])
+
+    def connect(self):
+        return self.session_provider.engine.connect()
+
+    def _save(self, datasets: list[Dataset]):
+        """Only do upserts. Never delete. Rows get only deleted when an entire Dataset is removed."""
+        datasets_entities = []
+        revision_entities = []
+        file_entities = []
+
+        for dataset in datasets:
+            datasets_entities.append(dataset.model_dump(exclude={"revisions"}))
+            for revision in dataset.revisions:
+                revision_entities.append(
+                    {
+                        **revision.model_dump(
+                            exclude={"is_squashed", "modified_files"}
+                        ),
+                        "dataset_id": dataset.dataset_id,
+                    }
+                )
+                for file in revision.modified_files:
+                    file_entities.append(
+                        {
+                            **file.model_dump(),
+                            "dataset_id": dataset.dataset_id,
+                            "revision_id": revision.revision_id,
+                        }
+                    )
+
+        with self.connect() as connection:
+            try:
+                self._upsert(connection, dataset_table, datasets_entities)
+                self._upsert(connection, revision_table, revision_entities)
+                self._upsert(connection, file_table, file_entities)
+            except Exception:
+                connection.rollback()
+                raise
+            else:
+                connection.commit()
 
     def destroy(self, dataset: Dataset):
-        self.session.delete(dataset)
-        self.session.commit()
+        with self.connect() as connection:
+            try:
+                # Delete modified files related to the dataset
+                file_table.delete().where(
+                    file_table.c.dataset_id == dataset.dataset_id
+                ).execute()
+
+                # Delete revisions related to the dataset
+                revision_table.delete().where(
+                    revision_table.c.dataset_id == dataset.dataset_id
+                ).execute()
+
+                # Delete the dataset itself
+                dataset_table.delete().where(
+                    dataset_table.c.dataset_id == dataset.dataset_id
+                ).execute()
+
+                connection.commit()
+            except Exception:
+                connection.rollback()
+                raise
 
     def next_identity(self):
         return str(uuid.uuid4())
+
+    # TODO: consider moving the IngestionJobSummary methods to a different Repository
+    def save_ingestion_job_summary(self, ingestion_job_summary: IngestionJobSummary):
+        ingestion_job_summary_entities = [
+            ingestion_job_summary.model_dump(exclude={"task_summaries"})
+        ]
+        task_summary_entities = []
+        for task_summary in ingestion_job_summary.task_summaries:
+            task_summary_entities.append(
+                {
+                    **task_summary.model_dump(),
+                    "ingestion_job_summary_id": ingestion_job_summary.ingestion_job_summary_id,
+                }
+            )
+
+        with self.session_provider.engine.connect() as connection:
+            try:
+                self._upsert(
+                    connection,
+                    ingestion_job_summary_table,
+                    ingestion_job_summary_entities,
+                )
+                if task_summary_entities:
+                    self._upsert(connection, task_summary_table, task_summary_entities)
+            except Exception:
+                connection.rollback()
+                raise
+            else:
+                connection.commit()
+
+    def load_ingestion_job_summaries(self) -> list[IngestionJobSummary]:
+        ingestion_job_summary_ids = [
+            row.ingestion_job_summary_id
+            for row in self.session.query(
+                ingestion_job_summary_table.c.ingestion_job_summary_id
+            )
+        ]
+
+        ingestion_job_summary_rows = list(
+            self.session.query(ingestion_job_summary_table).filter(
+                ingestion_job_summary_table.c.ingestion_job_summary_id.in_(
+                    ingestion_job_summary_ids
+                )
+            )
+        )
+
+        task_summary_entities_per_job_summary = {}
+        rows = (
+            self.session.query(task_summary_table)
+            .filter(
+                task_summary_table.c.ingestion_job_summary_id.in_(
+                    ingestion_job_summary_ids
+                )
+            )
+            .order_by(task_summary_table.c.ingestion_job_summary_id)
+        )
+
+        for ingestion_job_summary_id, task_summaries_rows in itertools.groupby(
+            rows, key=lambda row: row.ingestion_job_summary_id
+        ):
+            task_summary_entities_per_job_summary[ingestion_job_summary_id] = list(
+                task_summaries_rows
+            )
+
+        ingestion_job_summaries = []
+        for ingestion_job_summary_row in ingestion_job_summary_rows:
+            task_summaries = [
+                TaskSummary.model_validate(row)
+                for row in task_summary_entities_per_job_summary.get(
+                    ingestion_job_summary_row.ingestion_job_summary_id, []
+                )
+            ]
+
+            ingestion_job_summaries.append(
+                IngestionJobSummary.model_validate(
+                    {
+                        **ingestion_job_summary_row._mapping,
+                        "task_summaries": task_summaries,
+                    }
+                )
+            )
+        return ingestion_job_summaries