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

ingestify/__init__.py

@@ -8,4 +8,4 @@ if not __INGESTIFY_SETUP__:
     from .infra import retrieve_http
     from .source_base import Source, DatasetResource
 
-__version__ = "0.3.3"
+__version__ = "0.4.0"

ingestify/application/dataset_store.py

@@ -58,8 +58,7 @@ class DatasetStore:
             self.event_bus.dispatch(event)
 
     def save_ingestion_job_summary(self, ingestion_job_summary):
-        self.dataset_repository.session.add(ingestion_job_summary)
-        self.dataset_repository.session.commit()
+        self.dataset_repository.save_ingestion_job_summary(ingestion_job_summary)
 
     def get_dataset_collection(
         self,
@@ -271,6 +270,7 @@ class DatasetStore:
             metadata=metadata,
             created_at=now,
             updated_at=now,
+            last_modified_at=None,  # Not known at this moment
         )
         revision = self.add_revision(dataset, files, revision_source, description)
 
@@ -298,8 +298,8 @@ class DatasetStore:
                 )
 
             loaded_file = LoadedFile(
-                _stream=get_stream if lazy else get_stream(file),
-                **asdict(file),
+                stream_=get_stream if lazy else get_stream(file),
+                **file.model_dump(),
             )
             files[file.file_id] = loaded_file
         return FileCollection(files, auto_rewind=auto_rewind)

ingestify/application/ingestion_engine.py

@@ -21,8 +21,13 @@ class IngestionEngine:
     def add_ingestion_plan(self, ingestion_plan: IngestionPlan):
         self.loader.add_ingestion_plan(ingestion_plan)
 
-    def load(self, dry_run: bool = False, provider: Optional[str] = None):
-        self.loader.collect_and_run(dry_run=dry_run, provider=provider)
+    def load(
+        self,
+        dry_run: bool = False,
+        provider: Optional[str] = None,
+        source: Optional[str] = None,
+    ):
+        self.loader.collect_and_run(dry_run=dry_run, provider=provider, source=source)
 
     def list_datasets(self, as_count: bool = False):
         """Consider moving this to DataStore"""

ingestify/application/loader.py

@@ -29,7 +29,12 @@ class Loader:
     def add_ingestion_plan(self, ingestion_plan: IngestionPlan):
         self.ingestion_plans.append(ingestion_plan)
 
-    def collect_and_run(self, dry_run: bool = False, provider: Optional[str] = None):
+    def collect_and_run(
+        self,
+        dry_run: bool = False,
+        provider: Optional[str] = None,
+        source: Optional[str] = None,
+    ):
         # First collect all selectors, before discovering datasets
         selectors = {}
         for ingestion_plan in self.ingestion_plans:
@@ -42,6 +47,13 @@ class Loader:
                     )
                     continue
 
+            if source is not None:
+                if ingestion_plan.source.name != source:
+                    logger.info(
+                        f"Skipping {ingestion_plan} because source doesn't match '{source}'"
+                    )
+                    continue
+
             static_selectors = [
                 selector
                 for selector in ingestion_plan.selectors
@@ -60,6 +72,7 @@ class Loader:
 
                     # TODO: consider making this lazy and fetch once per Source instead of
                     #       once per IngestionPlan
+                    # TODO: Log exception when `discover_selectors` fails
                     all_selectors = ingestion_plan.source.discover_selectors(
                         ingestion_plan.dataset_type
                     )

ingestify/cmdline.py

@@ -58,7 +58,14 @@ def cli():
     help="bucket",
     type=str,
 )
-@click.option("--debug", "debug", required=False, help="Debugging enabled", type=bool)
+@click.option(
+    "--debug",
+    "debug",
+    required=False,
+    help="Debugging enabled",
+    is_flag=True,
+    type=bool,
+)
 @click.option(
     "--dry-run",
     "dry_run",
@@ -74,11 +81,19 @@ def cli():
     help="Provider - only run tasks for a single provider",
     type=str,
 )
+@click.option(
+    "--source",
+    "source",
+    required=False,
+    help="Source - only run tasks for a single source",
+    type=str,
+)
 def run(
     config_file: str,
     bucket: Optional[str],
     dry_run: Optional[bool],
     provider: Optional[str],
+    source: Optional[str],
     debug: Optional[bool],
 ):
     try:
@@ -90,7 +105,10 @@ def run(
             logger.exception(f"Failed due a configuration error: {e}")
             sys.exit(1)
 
-    engine.load(dry_run=dry_run, provider=provider)
+    if debug:
+        logging.getLogger("root").setLevel(logging.DEBUG)
+
+    engine.load(dry_run=dry_run, provider=provider, source=source)
 
     logger.info("Done")
 

ingestify/domain/models/base.py

@@ -1,22 +1,5 @@
-from functools import partial
-from typing import ClassVar, Any, Optional
-
-import pydantic
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict
 
 
-# class BaseModel(PydanticBaseModel):
-#     model_config = ConfigDict(arbitrary_types_allowed=True)
-#
-#     _sa_instance_state: Optional[dict] = None
-from sqlalchemy.orm import MappedAsDataclass
-
-
-class BaseModel(
-    MappedAsDataclass,
-    # DeclarativeBase,
-    dataclass_callable=partial(
-        pydantic.dataclasses.dataclass, config=ConfigDict(arbitrary_types_allowed=True)
-    ),
-):
-    pass
+class BaseModel(PydanticBaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)

ingestify/domain/models/dataset/collection_metadata.py

@@ -6,7 +6,8 @@ from typing import Optional
 @dataclass
 class DatasetCollectionMetadata:
     # This can be useful to figure out if a backfill is required
-    first_modified: Optional[datetime]
+    # TODO - Note: not stored at Dataset level and requires joined query to retrieve
+    # first_modified: Optional[datetime]
 
     # Use the last modified to only retrieve datasets that are changed
     last_modified: Optional[datetime]

ingestify/domain/models/dataset/dataset.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from enum import Enum
 from typing import List, Optional
-from pydantic import Field
+from pydantic import Field, field_validator
 
 from ingestify.utils import utcnow
 from .dataset_state import DatasetState
@@ -22,7 +22,17 @@ class Dataset(BaseModel):
     metadata: dict
     created_at: datetime
     updated_at: datetime
+
     revisions: List[Revision] = Field(default_factory=list)
+    # The last_modified_at is equal to the max modified_at of all files in all revisions
+    last_modified_at: Optional[datetime]
+
+    @field_validator("identifier", mode="before")
+    @classmethod
+    def parse_identifier(cls, value):
+        if not isinstance(value, Identifier):
+            return Identifier(value)
+        return value
 
     @property
     def is_complete(self):
@@ -35,6 +45,13 @@ class Dataset(BaseModel):
         self.revisions.append(revision)
         self.updated_at = utcnow()
 
+        if self.last_modified_at:
+            self.last_modified_at = max(
+                self.last_modified_at, revision.last_modified_at
+            )
+        else:
+            self.last_modified_at = revision.last_modified_at
+
     def update_metadata(self, name: str, metadata: dict, state: DatasetState) -> bool:
         changed = False
         if self.name != name:

ingestify/domain/models/dataset/file.py

@@ -116,18 +116,18 @@ class LoadedFile(BaseModel):
     data_serialization_format: Optional[str]  # Example: 'json'
     storage_compression_method: Optional[str]  # Example: 'gzip'
     storage_path: Path
-    _stream: Union[BinaryIO, BytesIO, Callable[[], Awaitable[Union[BinaryIO, BytesIO]]]]
+    stream_: Union[BinaryIO, BytesIO, Callable[[], Awaitable[Union[BinaryIO, BytesIO]]]]
     revision_id: Optional[int] = None  # This can be used when a Revision is squashed
 
     def load_stream(self):
-        if callable(self._stream):
-            self._stream = self._stream(self)
+        if callable(self.stream_):
+            self.stream_ = self.stream_(self)
 
     @property
     def stream(self):
-        if callable(self._stream):
+        if callable(self.stream_):
             raise Exception("You should load the stream first using `load_stream`")
-        return self._stream
+        return self.stream_
 
 
 __all__ = ["File", "DraftFile", "LoadedFile"]

ingestify/domain/models/dataset/revision.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from enum import Enum
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 from typing_extensions import TypedDict
 
@@ -32,10 +32,14 @@ class Revision(BaseModel):
     created_at: datetime
     description: str
     modified_files: List[File]
-    source: RevisionSource
+    source: Optional[RevisionSource]
     is_squashed: bool = False
     state: RevisionState = RevisionState.PENDING_VALIDATION
 
+    @property
+    def last_modified_at(self):
+        return max(file.modified_at for file in self.modified_files)
+
     @property
     def modified_files_map(self) -> Dict[str, File]:
         return {file.file_id: file for file in self.modified_files}

ingestify/domain/models/ingestion/ingestion_job.py

@@ -2,6 +2,7 @@ import itertools
 import json
 import logging
 import uuid
+from enum import Enum
 from typing import Optional, Iterator
 
 from ingestify import retrieve_http
@@ -17,6 +18,7 @@ from ingestify.domain.models.resources.dataset_resource import (
     DatasetResource,
 )
 from ingestify.domain.models.task.task_summary import TaskSummary
+from ingestify.exceptions import SaveError
 from ingestify.utils import TaskExecutor, chunker
 
 logger = logging.getLogger(__name__)
@@ -120,21 +122,27 @@ class UpdateDatasetTask(Task):
         with TaskSummary.update(
             self.task_id, dataset_identifier=dataset_identifier
         ) as task_summary:
-            revision = self.store.update_dataset(
-                dataset=self.dataset,
-                name=self.dataset_resource.name,
-                state=self.dataset_resource.state,
-                metadata=self.dataset_resource.metadata,
-                files={
-                    file_id: task_summary.record_load_file(
-                        lambda: load_file(file_resource, dataset=self.dataset),
-                        metadata={"file_id": file_id},
-                    )
-                    for file_id, file_resource in self.dataset_resource.files.items()
-                },
-                revision_source=revision_source,
-            )
-            task_summary.set_stats_from_revision(revision)
+
+            files = {
+                file_id: task_summary.record_load_file(
+                    lambda: load_file(file_resource, dataset=self.dataset),
+                    metadata={"file_id": file_id},
+                )
+                for file_id, file_resource in self.dataset_resource.files.items()
+            }
+
+            try:
+                revision = self.store.update_dataset(
+                    dataset=self.dataset,
+                    name=self.dataset_resource.name,
+                    state=self.dataset_resource.state,
+                    metadata=self.dataset_resource.metadata,
+                    files=files,
+                    revision_source=revision_source,
+                )
+                task_summary.set_stats_from_revision(revision)
+            except Exception as e:
+                raise SaveError("Could not update dataset") from e
 
         return task_summary
 
@@ -159,24 +167,28 @@ class CreateDatasetTask(Task):
         )
 
         with TaskSummary.create(self.task_id, dataset_identifier) as task_summary:
-            revision = self.store.create_dataset(
-                dataset_type=self.dataset_resource.dataset_type,
-                provider=self.dataset_resource.provider,
-                dataset_identifier=dataset_identifier,
-                name=self.dataset_resource.name,
-                state=self.dataset_resource.state,
-                metadata=self.dataset_resource.metadata,
-                files={
-                    file_id: task_summary.record_load_file(
-                        lambda: load_file(file_resource, dataset=None),
-                        metadata={"file_id": file_id},
-                    )
-                    for file_id, file_resource in self.dataset_resource.files.items()
-                },
-                revision_source=revision_source,
-            )
+            files = {
+                file_id: task_summary.record_load_file(
+                    lambda: load_file(file_resource, dataset=None),
+                    metadata={"file_id": file_id},
+                )
+                for file_id, file_resource in self.dataset_resource.files.items()
+            }
+            try:
+                revision = self.store.create_dataset(
+                    dataset_type=self.dataset_resource.dataset_type,
+                    provider=self.dataset_resource.provider,
+                    dataset_identifier=dataset_identifier,
+                    name=self.dataset_resource.name,
+                    state=self.dataset_resource.state,
+                    metadata=self.dataset_resource.metadata,
+                    files=files,
+                    revision_source=revision_source,
+                )
 
-            task_summary.set_stats_from_revision(revision)
+                task_summary.set_stats_from_revision(revision)
+            except Exception as e:
+                raise SaveError("Could not create dataset") from e
 
         return task_summary
 
@@ -209,6 +221,7 @@ class IngestionJob:
         with ingestion_job_summary.record_timing("get_dataset_collection"):
             dataset_collection_metadata = store.get_dataset_collection(
                 dataset_type=self.ingestion_plan.dataset_type,
+                provider=self.ingestion_plan.source.provider,
                 data_spec_versions=self.selector.data_spec_versions,
                 selector=self.selector,
                 metadata_only=True,
@@ -218,27 +231,38 @@ class IngestionJob:
         # There are two different, but similar flows here:
         # 1. The discover_datasets returns a list, and the entire list can be processed at once
         # 2. The discover_datasets returns an iterator of batches, in this case we need to process each batch
-        with ingestion_job_summary.record_timing("find_datasets"):
-            # Timing might be incorrect as it is an iterator
-            dataset_resources = self.ingestion_plan.source.find_datasets(
-                dataset_type=self.ingestion_plan.dataset_type,
-                data_spec_versions=self.selector.data_spec_versions,
-                dataset_collection_metadata=dataset_collection_metadata,
-                **self.selector.custom_attributes,
-            )
+        try:
+            with ingestion_job_summary.record_timing("find_datasets"):
+                dataset_resources = self.ingestion_plan.source.find_datasets(
+                    dataset_type=self.ingestion_plan.dataset_type,
+                    data_spec_versions=self.selector.data_spec_versions,
+                    dataset_collection_metadata=dataset_collection_metadata,
+                    **self.selector.custom_attributes,
+                )
 
-        finish_task_timer = ingestion_job_summary.start_timing("tasks")
+                # We need to include the to_batches as that will start the generator
+                batches = to_batches(dataset_resources)
+        except Exception as e:
+            logger.exception("Failed to find datasets")
 
-        batches = to_batches(dataset_resources)
+            ingestion_job_summary.set_exception(e)
+            yield ingestion_job_summary
+            return
+
+        finish_task_timer = ingestion_job_summary.start_timing("tasks")
 
         while True:
             try:
                 batch = next(batches)
             except StopIteration:
                 break
-            except Exception:
-                # TODO: handle exception on IngestionJob level
-                raise
+            except Exception as e:
+                logger.exception("Failed to fetch next batch")
+
+                finish_task_timer()
+                ingestion_job_summary.set_exception(e)
+                yield ingestion_job_summary
+                return
 
             dataset_identifiers = [
                 Identifier.create_from_selector(

ingestify/domain/models/ingestion/ingestion_job_summary.py

@@ -1,24 +1,31 @@
 import uuid
 from contextlib import contextmanager
 from datetime import datetime, timedelta
+from enum import Enum
 from typing import Optional, List, TYPE_CHECKING
 from pydantic import Field
 
 from ingestify.domain import Selector, DataSpecVersionCollection
 from ingestify.domain.models.base import BaseModel
-from ingestify.domain.models.task.task_summary import TaskSummary, TaskStatus
+from ingestify.domain.models.task.task_summary import TaskSummary, TaskState
 from ingestify.domain.models.timing import Timing
-from ingestify.utils import utcnow
+from ingestify.utils import utcnow, HasTiming
 
 if TYPE_CHECKING:
     from ingestify.domain.models.ingestion.ingestion_job import IngestionJob
 
 
+class IngestionJobState(str, Enum):
+    RUNNING = "RUNNING"
+    FINISHED = "FINISHED"
+    FAILED = "FAILED"
+
+
 def format_duration(duration: timedelta):
     return f"{duration.total_seconds():.2f}sec"
 
 
-class IngestionJobSummary(BaseModel):
+class IngestionJobSummary(BaseModel, HasTiming):
     ingestion_job_summary_id: str
     ingestion_job_id: str
 
@@ -30,8 +37,8 @@ class IngestionJobSummary(BaseModel):
     selector: Selector
 
     started_at: datetime = Field(default_factory=utcnow)
-    finished_at: Optional[datetime] = None
-    timings: List[Timing] = Field(default_factory=list)
+    ended_at: Optional[datetime] = None
+    state: IngestionJobState = IngestionJobState.RUNNING
     task_summaries: List[TaskSummary] = Field(default_factory=list)
 
     skipped_datasets: int = 0
@@ -52,20 +59,6 @@ class IngestionJobSummary(BaseModel):
         )
         return cls(**args)
 
-    @contextmanager
-    def record_timing(self, name: str):
-        start = utcnow()
-        yield
-        self.timings.append(Timing(name=name, started_at=start, ended_at=utcnow()))
-
-    def start_timing(self, name):
-        start = utcnow()
-
-        def finish():
-            self.timings.append(Timing(name=name, started_at=start, ended_at=utcnow()))
-
-        return finish
-
     def add_task_summaries(self, task_summaries: List[TaskSummary]):
         self.task_summaries.extend(task_summaries)
 
@@ -75,46 +68,61 @@ class IngestionJobSummary(BaseModel):
     def task_count(self):
         return len(self.task_summaries)
 
-    def set_finished(self):
+    def _set_ended(self):
         self.failed_tasks = len(
-            [task for task in self.task_summaries if task.status == TaskStatus.FAILED]
+            [task for task in self.task_summaries if task.state == TaskState.FAILED]
         )
         self.successful_tasks = len(
-            [task for task in self.task_summaries if task.status == TaskStatus.FINISHED]
+            [task for task in self.task_summaries if task.state == TaskState.FINISHED]
         )
         self.ignored_successful_tasks = len(
             [
                 task
                 for task in self.task_summaries
-                if task.status == TaskStatus.FINISHED_IGNORED
+                if task.state == TaskState.FINISHED_IGNORED
             ]
         )
-        self.finished_at = utcnow()
+        self.ended_at = utcnow()
+
+        # Only keep failed tasks. Rest isn't interesting
+        self.task_summaries = [
+            task for task in self.task_summaries if task.state == TaskState.FAILED
+        ]
+
+    def set_finished(self):
+        self.state = IngestionJobState.FINISHED
+        self._set_ended()
+
+    def set_exception(self, e: Exception):
+        self.state = IngestionJobState.FAILED
+        self._set_ended()
 
     @property
     def duration(self) -> timedelta:
-        return self.finished_at - self.started_at
+        return self.ended_at - self.started_at
 
     def output_report(self):
-        print(f"\nIngestionJobSummary finished in {format_duration(self.duration)}")
-        print("--------------------")
-        print(f"  - IngestionPlan:")
-        print(f"        Source: {self.source_name}")
-        print(f"        Provider: {self.provider}")
-        print(f"        DatasetType: {self.dataset_type}")
-        print(f"  - Selector: {self.selector}")
-        print(f"  - Timings: ")
+        print(
+            f"\nIngestionJobSummary {self.state.value} in {format_duration(self.duration)}"
+        )
+        print("********************************")
+        print(f"*  - IngestionPlan:")
+        print(f"*        Source: {self.source_name}")
+        print(f"*        Provider: {self.provider}")
+        print(f"*        DatasetType: {self.dataset_type}")
+        print(f"*  - Selector: {self.selector}")
+        print(f"*  - Timings: ")
         for timing in self.timings:
-            print(f"    - {timing.name}: {format_duration(timing.duration)}")
+            print(f"*    - {timing.name}: {format_duration(timing.duration)}")
         print(
-            f"  - Tasks: {len(self.task_summaries)} - {(len(self.task_summaries) / self.duration.total_seconds()):.1f} tasks/sec"
+            f"*  - Tasks: {len(self.task_summaries)} - {(len(self.task_summaries) / self.duration.total_seconds()):.1f} tasks/sec"
         )
 
-        print(f"    - Failed tasks: {self.failed_tasks}")
-        print(f"    - Successful tasks: {self.successful_tasks}")
-        print(f"    - Successful ignored tasks: {self.ignored_successful_tasks}")
-        print(f"    - Skipped datasets: {self.skipped_datasets}")
-        print("--------------------")
+        print(f"*    - Failed tasks: {self.failed_tasks}")
+        print(f"*    - Successful tasks: {self.successful_tasks}")
+        print(f"*    - Successful ignored tasks: {self.ignored_successful_tasks}")
+        print(f"*    - Skipped datasets: {self.skipped_datasets}")
+        print("********************************")
 
     def __enter__(self):
         return self

ingestify/domain/models/task/task_summary.py

@@ -10,13 +10,12 @@ from ingestify.domain.models.base import BaseModel
 from ingestify.domain.models.dataset.identifier import Identifier
 from ingestify.domain.models.timing import Timing
 from ingestify.exceptions import IngestifyError
-from ingestify.utils import utcnow
-
+from ingestify.utils import utcnow, HasTiming
 
 logger = logging.getLogger(__name__)
 
 
-class TaskStatus(str, Enum):
+class TaskState(str, Enum):
     RUNNING = "RUNNING"
     FINISHED = "FINISHED"
     FINISHED_IGNORED = "FINISHED_IGNORED"  # Finished, but didn't produce any new data
@@ -28,7 +27,7 @@ class Operation(str, Enum):
     UPDATE = "UPDATE"
 
 
-class TaskSummary(BaseModel):
+class TaskSummary(BaseModel, HasTiming):
     task_id: str
     started_at: datetime
     operation: Operation
@@ -37,8 +36,7 @@ class TaskSummary(BaseModel):
     persisted_file_count: int = 0
     bytes_retrieved: int = 0
     last_modified: Optional[datetime] = None
-    status: TaskStatus = TaskStatus.RUNNING
-    timings: List[Timing] = Field(default_factory=list)
+    state: TaskState = TaskState.RUNNING
 
     @field_validator("dataset_identifier", mode="before")
     @classmethod
@@ -48,27 +46,8 @@ class TaskSummary(BaseModel):
         return value
 
     def record_load_file(self, fn, metadata: dict):
-        start = utcnow()
-        try:
-            result = None
+        with self.record_timing(f"Load of {metadata.get('file_id', 'file')}", metadata):
             return fn()
-        except Exception as e:
-            result = {
-                "type": type(e).__name__,
-                "message": str(e),
-                "traceback": traceback.format_exc(),
-            }
-            raise e
-        finally:
-            metadata = dict(result=result, **metadata)
-            self.timings.append(
-                Timing(
-                    name=f"Load of {metadata.get('file_id', 'file')}",
-                    started_at=start,
-                    ended_at=utcnow(),
-                    metadata=metadata,
-                )
-            )
 
     @classmethod
     @contextmanager
@@ -83,10 +62,10 @@ class TaskSummary(BaseModel):
         try:
             yield task_summary
 
-            task_summary.set_status(TaskStatus.FINISHED)
+            task_summary.set_state(TaskState.FINISHED)
         except Exception as e:
             logger.exception(f"Failed to execute task.")
-            task_summary.set_status(TaskStatus.FAILED)
+            task_summary.set_state(TaskState.FAILED)
 
             # When the error comes from our own code, make sure it will be raised to the highest level
             # raise
@@ -111,8 +90,8 @@ class TaskSummary(BaseModel):
                 file.modified_at for file in revision.modified_files
             )
         else:
-            self.status = TaskStatus.FINISHED_IGNORED
+            self.state = TaskState.FINISHED_IGNORED
 
-    def set_status(self, status: TaskStatus):
-        if self.status == TaskStatus.RUNNING:
-            self.status = status
+    def set_state(self, state: TaskState):
+        if self.state == TaskState.RUNNING:
+            self.state = state