ingestify 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

ingestify/domain/models/ingestion/ingestion_job.py

@@ -0,0 +1,292 @@
+import itertools
+import json
+import logging
+import uuid
+from typing import Optional
+
+from ingestify import retrieve_http
+from ingestify.application.dataset_store import DatasetStore
+from ingestify.domain import Selector, Identifier, TaskSet, Dataset, DraftFile, Task
+from ingestify.domain.models.dataset.revision import RevisionSource, SourceType
+from ingestify.domain.models.ingestion.ingestion_job_summary import (
+    IngestionJobSummary,
+)
+from ingestify.domain.models.ingestion.ingestion_plan import IngestionPlan
+from ingestify.domain.models.resources.dataset_resource import (
+    FileResource,
+    DatasetResource,
+)
+from ingestify.domain.models.task.task_summary import TaskSummary
+from ingestify.utils import TaskExecutor, chunker
+
+logger = logging.getLogger(__name__)
+
+
+DEFAULT_CHUNK_SIZE = 1000
+
+
+def run_task(task):
+    logger.info(f"Running task {task}")
+    return task.run()
+
+
+def to_batches(input_):
+    if isinstance(input_, list):
+        batches = [input_]
+    else:
+        # Assume it's an iterator. Peek what's inside, and put it back
+        try:
+            peek = next(input_)
+        except StopIteration:
+            # Nothing to batch
+            return []
+
+        input_ = itertools.chain([peek], input_)
+
+        if not isinstance(peek, list):
+            batches = chunker(input_, DEFAULT_CHUNK_SIZE)
+        else:
+            batches = input_
+    return batches
+
+
+def load_file(
+    file_resource: FileResource, dataset: Optional[Dataset] = None
+) -> Optional[DraftFile]:
+    current_file = None
+    if dataset:
+        current_file = dataset.current_revision.modified_files_map.get(
+            file_resource.file_id
+        )
+
+    if file_resource.json_content is not None:
+        # Empty dictionary is allowed
+        file = DraftFile.from_input(
+            file_=json.dumps(file_resource.json_content, indent=4),
+            data_serialization_format="json",
+            data_feed_key=file_resource.data_feed_key,
+            data_spec_version=file_resource.data_spec_version,
+            modified_at=file_resource.last_modified,
+        )
+        if current_file and current_file.tag == file.tag:
+            # Nothing changed
+            return None
+        return file
+    elif file_resource.url:
+        http_options = {}
+        if file_resource.http_options:
+            for k, v in file_resource.http_options.items():
+                http_options[f"http_{k}"] = v
+
+        return retrieve_http(
+            url=file_resource.url,
+            current_file=current_file,
+            file_data_feed_key=file_resource.data_feed_key,
+            file_data_spec_version=file_resource.data_spec_version,
+            file_data_serialization_format=file_resource.data_serialization_format
+            or "txt",
+            last_modified=file_resource.last_modified,
+            **http_options,
+            **file_resource.loader_kwargs,
+        )
+    else:
+        return file_resource.file_loader(
+            file_resource,
+            current_file,
+            # TODO: check how to fix this with typehints
+            **file_resource.loader_kwargs,
+        )
+
+
+class UpdateDatasetTask(Task):
+    def __init__(
+        self,
+        dataset: Dataset,
+        dataset_resource: DatasetResource,
+        store: DatasetStore,
+    ):
+        self.dataset = dataset
+        self.dataset_resource = dataset_resource
+        self.store = store
+        self.task_id = str(uuid.uuid1())
+
+    def run(self):
+        dataset_identifier = Identifier(**self.dataset_resource.dataset_resource_id)
+
+        revision_source = RevisionSource(
+            source_id=self.task_id, source_type=SourceType.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)
+
+        return task_summary
+
+    def __repr__(self):
+        return f"UpdateDatasetTask({self.dataset_resource.provider} -> {self.dataset_resource.dataset_resource_id})"
+
+
+class CreateDatasetTask(Task):
+    def __init__(
+        self,
+        dataset_resource: DatasetResource,
+        store: DatasetStore,
+    ):
+        self.dataset_resource = dataset_resource
+        self.store = store
+        self.task_id = str(uuid.uuid1())
+
+    def run(self):
+        dataset_identifier = Identifier(**self.dataset_resource.dataset_resource_id)
+        revision_source = RevisionSource(
+            source_id=self.task_id, source_type=SourceType.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,
+            )
+
+            task_summary.set_stats_from_revision(revision)
+
+        return task_summary
+
+    def __repr__(self):
+        return f"CreateDatasetTask({self.dataset_resource.provider} -> {self.dataset_resource.dataset_resource_id})"
+
+
+class IngestionJob:
+    def __init__(
+        self,
+        ingestion_job_id: str,
+        ingestion_plan: IngestionPlan,
+        selector: Selector,
+    ):
+        self.ingestion_job_id = ingestion_job_id
+        self.ingestion_plan = ingestion_plan
+        self.selector = selector
+
+    def execute(
+        self, store: DatasetStore, task_executor: TaskExecutor
+    ) -> IngestionJobSummary:
+        with IngestionJobSummary.new(ingestion_job=self) as ingestion_job_summary:
+            with ingestion_job_summary.record_timing("get_dataset_collection"):
+                dataset_collection_metadata = store.get_dataset_collection(
+                    dataset_type=self.ingestion_plan.dataset_type,
+                    data_spec_versions=self.selector.data_spec_versions,
+                    selector=self.selector,
+                    metadata_only=True,
+                ).metadata
+
+            # 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
+                datasets = 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,
+                )
+
+            batches = to_batches(datasets)
+
+            with ingestion_job_summary.record_timing("tasks"):
+                for batch in batches:
+                    dataset_identifiers = [
+                        Identifier.create_from_selector(
+                            self.selector, **dataset_resource.dataset_resource_id
+                        )
+                        # We have to pass the data_spec_versions here as a Source can add some
+                        # extra data to the identifier which is retrieved in a certain data format
+                        for dataset_resource in batch
+                    ]
+
+                    # Load all available datasets based on the discovered dataset identifiers
+                    dataset_collection = store.get_dataset_collection(
+                        dataset_type=self.ingestion_plan.dataset_type,
+                        # Assume all DatasetResources share the same provider
+                        provider=batch[0].provider,
+                        selector=dataset_identifiers,
+                    )
+
+                    skip_count = 0
+
+                    task_set = TaskSet()
+                    for dataset_resource in batch:
+                        dataset_identifier = Identifier.create_from_selector(
+                            self.selector, **dataset_resource.dataset_resource_id
+                        )
+
+                        if dataset := dataset_collection.get(dataset_identifier):
+                            if self.ingestion_plan.fetch_policy.should_refetch(
+                                dataset, dataset_resource
+                            ):
+                                task_set.add(
+                                    UpdateDatasetTask(
+                                        dataset=dataset,  # Current dataset from the database
+                                        dataset_resource=dataset_resource,  # Most recent dataset_resource
+                                        store=store,
+                                    )
+                                )
+                            else:
+                                skip_count += 1
+                        else:
+                            if self.ingestion_plan.fetch_policy.should_fetch(
+                                dataset_resource
+                            ):
+                                task_set.add(
+                                    CreateDatasetTask(
+                                        dataset_resource=dataset_resource,
+                                        store=store,
+                                    )
+                                )
+                            else:
+                                skip_count += 1
+
+                    if task_set:
+                        logger.info(
+                            f"Discovered {len(dataset_identifiers)} datasets from {self.ingestion_plan.source.__class__.__name__} "
+                            f"using selector {self.selector} => {len(task_set)} tasks. {skip_count} skipped."
+                        )
+                        logger.info(f"Running {len(task_set)} tasks")
+                        ingestion_job_summary.add_task_summaries(
+                            task_executor.run(run_task, task_set)
+                        )
+                    else:
+                        logger.info(
+                            f"Discovered {len(dataset_identifiers)} datasets from {self.ingestion_plan.source.__class__.__name__} "
+                            f"using selector {self.selector} => nothing to do"
+                        )
+
+        return ingestion_job_summary

ingestify/domain/models/ingestion/ingestion_job_summary.py

@@ -0,0 +1,106 @@
+from contextlib import contextmanager
+from datetime import datetime, timedelta
+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.timing import Timing
+from ingestify.utils import utcnow
+
+if TYPE_CHECKING:
+    from ingestify.domain.models.ingestion.ingestion_job import IngestionJob
+
+
+def format_duration(duration: timedelta):
+    return f"{duration.total_seconds():.2}sec"
+
+
+class IngestionJobSummary(BaseModel):
+    ingestion_job_id: str
+
+    # From the IngestionPlan
+    source_name: str
+    dataset_type: str
+    data_spec_versions: DataSpecVersionCollection
+    selector: Selector
+
+    started_at: datetime = Field(default_factory=utcnow)
+    finished_at: Optional[datetime] = None
+    timings: List[Timing] = Field(default_factory=list)
+    task_summaries: List[TaskSummary] = Field(default_factory=list)
+
+    failed_tasks: int = 0
+    successful_tasks: int = 0
+    ignored_successful_tasks: int = 0
+
+    @classmethod
+    def new(cls, ingestion_job: "IngestionJob"):
+        args = dict(
+            ingestion_job_id=ingestion_job.ingestion_job_id,
+            source_name=ingestion_job.ingestion_plan.source.name,
+            dataset_type=ingestion_job.ingestion_plan.dataset_type,
+            data_spec_versions=ingestion_job.ingestion_plan.data_spec_versions,
+            selector=ingestion_job.selector,
+        )
+        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 add_task_summaries(self, task_summaries: List[TaskSummary]):
+        self.task_summaries.extend(task_summaries)
+
+    def set_finished(self):
+        self.failed_tasks = len(
+            [task for task in self.task_summaries if task.status == TaskStatus.FAILED]
+        )
+        self.successful_tasks = len(
+            [task for task in self.task_summaries if task.status == TaskStatus.FINISHED]
+        )
+        self.ignored_successful_tasks = len(
+            [
+                task
+                for task in self.task_summaries
+                if task.status == TaskStatus.FINISHED_IGNORED
+            ]
+        )
+        self.finished_at = utcnow()
+
+    @property
+    def duration(self) -> timedelta:
+        return self.finished_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"        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"  - Tasks: {len(self.task_summaries)} - {(len(self.task_summaries) / self.duration.total_seconds()):.1f} tasks/sec"
+        )
+
+        for status in [
+            TaskStatus.FAILED,
+            TaskStatus.FINISHED,
+            TaskStatus.FINISHED_IGNORED,
+        ]:
+            print(
+                f"    - {status.value.lower()}: {len([task for task in self.task_summaries if task.status == status])}"
+            )
+        print("--------------------")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass

ingestify/domain/models/{extract_job.py → ingestion/ingestion_plan.py}

@@ -1,15 +1,15 @@
-from dataclasses import dataclass
 from typing import List
 
 from ingestify.domain.models import Source, Selector
+from ingestify.domain.models.base import BaseModel
 from ingestify.domain.models.data_spec_version_collection import (
     DataSpecVersionCollection,
 )
 from ingestify.domain.models.fetch_policy import FetchPolicy
 
 
-@dataclass
-class ExtractJob:
+class IngestionPlan(BaseModel):
+
     source: Source
     selectors: List[Selector]
     fetch_policy: FetchPolicy
@@ -17,7 +17,7 @@ class ExtractJob:
     data_spec_versions: DataSpecVersionCollection
 
     def __repr__(self):
-        return f'<ExtractJob source="{self.source.name}" dataset_type="{self.dataset_type}">'
+        return f'<IngestionPlan source="{self.source.name}" dataset_type="{self.dataset_type}">'
 
     def __str__(self):
         return repr(self)

ingestify/domain/models/resources/dataset_resource.py

@@ -1,35 +1,40 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional, Callable, TYPE_CHECKING
+from typing import Optional, Callable, Any, Protocol, TYPE_CHECKING  # noqa
+from pydantic import Field
 
+from ingestify.domain.models.base import BaseModel
+from ingestify.domain.models.dataset.dataset_state import DatasetState
 from ingestify.exceptions import DuplicateFile
 
-if TYPE_CHECKING:
-    from ingestify.domain import DraftFile, File
-    from ingestify.domain.models.dataset.dataset import DatasetState
+from ingestify.domain.models import File, DraftFile
 
 
-@dataclass(frozen=True)
-class FileResource:
+class FileLoaderProtocol(Protocol):
+    def __call__(
+        self,
+        file_resource: "FileResource",
+        file: Optional["File"] = None,
+        **kwargs: Any,
+    ) -> Optional["DraftFile"]:
+        ...
+
+
+class FileResource(BaseModel):
     dataset_resource: "DatasetResource"
     file_id: str
     last_modified: datetime
     data_feed_key: str
     data_spec_version: str
-
-    # DataSerializationFormat is "json" in case of json_content, otherwise file_loader will return it
-    # data_serialization_format: str
-
     json_content: Optional[dict] = None
-
     url: Optional[str] = None
     http_options: Optional[dict] = None
+    # DataSerializationFormat is "json" in case of json_content, otherwise file_loader will return it
     data_serialization_format: Optional[str] = None
-
     file_loader: Optional[
         Callable[["FileResource", Optional["File"]], Optional["DraftFile"]]
     ] = None
-    loader_kwargs: dict = field(default_factory=dict)
+    loader_kwargs: dict = Field(default_factory=dict)
 
     def __post_init__(self):
         if self.json_content is None and not self.url and not self.file_loader:
@@ -38,27 +43,14 @@ class FileResource:
             )
 
 
-class DatasetResource:
-    def __init__(
-        self,
-        dataset_resource_id: dict,
-        /,
-        dataset_type: str,
-        provider: str,
-        name: str,
-        metadata: Optional[dict] = None,
-        state: Optional["DatasetState"] = None,
-    ):
-        from ingestify.domain.models.dataset.dataset import DatasetState
-
-        self.dataset_type = dataset_type
-        self.provider = provider
-        self.dataset_resource_id = dataset_resource_id
-        self.name = name
-        self.metadata = metadata or {}
-        self.state = state or DatasetState.COMPLETE
-
-        self.files = {}
+class DatasetResource(BaseModel):
+    dataset_resource_id: dict
+    dataset_type: str
+    provider: str
+    name: str
+    metadata: dict = Field(default_factory=dict)
+    state: DatasetState = Field(default_factory=lambda: DatasetState.COMPLETE)
+    files: dict[str, FileResource] = Field(default_factory=dict)
 
     def add_file(
         self,
@@ -72,8 +64,8 @@ class DatasetResource:
         data_serialization_format: Optional[str] = None,
         file_loader: Optional[
             Callable[
-                ["FileResource", Optional["File"]],
-                Optional["DraftFile"],
+                [FileResource, Optional[File]],
+                Optional[DraftFile],
             ]
         ] = None,
         loader_kwargs: Optional[dict] = None,

ingestify/domain/models/sink.py

@@ -1,16 +1,9 @@
 from abc import ABC, abstractmethod
 
-from ingestify.utils import ComponentFactory, ComponentRegistry
-
 from .dataset import Dataset
 
-sink_registry = ComponentRegistry()
-
 
-class Sink(ABC, metaclass=sink_registry.metaclass):
+class Sink(ABC):
     @abstractmethod
     def upsert(self, dataset: Dataset, data, params: dict):
         pass
-
-
-sink_factory = ComponentFactory.build_factory(Sink, sink_registry)

ingestify/domain/models/task/task.py

@@ -1,7 +1,9 @@
 from abc import ABC, abstractmethod
 
+from .task_summary import TaskSummary
+
 
 class Task(ABC):
     @abstractmethod
-    def run(self):
+    def run(self) -> TaskSummary:
         pass

ingestify/domain/models/task/task_summary.py

@@ -0,0 +1,118 @@
+import logging
+import traceback
+from contextlib import contextmanager
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+from pydantic import Field, field_validator
+
+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
+
+
+logger = logging.getLogger(__name__)
+
+
+class TaskStatus(str, Enum):
+    RUNNING = "RUNNING"
+    FINISHED = "FINISHED"
+    FINISHED_IGNORED = "FINISHED_IGNORED"  # Finished, but didn't produce any new data
+    FAILED = "FAILED"
+
+
+class Operation(str, Enum):
+    CREATE = "CREATE"
+    UPDATE = "UPDATE"
+
+
+class TaskSummary(BaseModel):
+    task_id: str
+    started_at: datetime
+    operation: Operation
+    dataset_identifier: Identifier
+    ended_at: Optional[datetime] = None
+    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)
+
+    @field_validator("dataset_identifier", mode="before")
+    @classmethod
+    def ensure_list(cls, value) -> Identifier:
+        if not isinstance(value, Identifier):
+            return Identifier(**value)
+        return value
+
+    def record_load_file(self, fn, metadata: dict):
+        start = utcnow()
+        try:
+            result = None
+            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
+    def new(cls, task_id: str, operation: Operation, dataset_identifier: Identifier):
+        start = utcnow()
+        task_summary = cls(
+            task_id=task_id,
+            started_at=start,
+            operation=operation,
+            dataset_identifier=dataset_identifier,
+        )
+        try:
+            yield task_summary
+
+            task_summary.set_status(TaskStatus.FINISHED)
+        except Exception as e:
+            logger.exception(f"Failed to execute task.")
+            task_summary.set_status(TaskStatus.FAILED)
+
+            # When the error comes from our own code, make sure it will be raised to the highest level
+            # raise
+            if isinstance(e, IngestifyError):
+                raise
+        finally:
+            task_summary.ended_at = utcnow()
+
+    @classmethod
+    def update(cls, task_id: str, dataset_identifier: Identifier):
+        return cls.new(task_id, Operation.UPDATE, dataset_identifier)
+
+    @classmethod
+    def create(cls, task_id: str, dataset_identifier: Identifier):
+        return cls.new(task_id, Operation.CREATE, dataset_identifier)
+
+    def set_stats_from_revision(self, revision: Optional["Revision"]):
+        if revision:
+            self.persisted_file_count = len(revision.modified_files)
+            self.bytes_retrieved = sum(file.size for file in revision.modified_files)
+            self.last_modified = max(
+                file.modified_at for file in revision.modified_files
+            )
+        else:
+            self.status = TaskStatus.FINISHED_IGNORED
+
+    def set_status(self, status: TaskStatus):
+        if self.status == TaskStatus.RUNNING:
+            self.status = status

ingestify/domain/models/timing.py

@@ -0,0 +1,16 @@
+from datetime import datetime
+from typing import Optional, Any
+from pydantic import BaseModel, ConfigDict
+
+
+class Timing(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    started_at: datetime
+    ended_at: datetime
+    metadata: Optional[dict[str, Any]] = None
+
+    @property
+    def duration(self):
+        return self.ended_at - self.started_at