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

ingestify/domain/models/ingestion/ingestion_job.py

@@ -0,0 +1,325 @@
+import itertools
+import json
+import logging
+import uuid
+from typing import Optional, Iterator
+
+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 = 1_000
+
+
+def run_task(task):
+    logger.info(f"Running task {task}")
+    return task.run()
+
+
+def to_batches(input_):
+    if isinstance(input_, list):
+        batches = iter(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 iter([])
+
+        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})"
+
+
+MAX_TASKS_PER_CHUNK = 10_000
+
+
+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
+    ) -> Iterator[IngestionJobSummary]:
+        is_first_chunk = True
+        ingestion_job_summary = IngestionJobSummary.new(ingestion_job=self)
+        # Process all items in batches. Yield a IngestionJobSummary per batch
+
+        logger.info("Finding metadata")
+        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
+        logger.info(f"Done: {dataset_collection_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
+            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")
+
+        batches = to_batches(dataset_resources)
+
+        while True:
+            try:
+                batch = next(batches)
+            except StopIteration:
+                break
+            except Exception:
+                # TODO: handle exception on IngestionJob level
+                raise
+
+            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,
+            )
+
+            skipped_datasets = 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:
+                        skipped_datasets += 1
+                else:
+                    if self.ingestion_plan.fetch_policy.should_fetch(dataset_resource):
+                        task_set.add(
+                            CreateDatasetTask(
+                                dataset_resource=dataset_resource,
+                                store=store,
+                            )
+                        )
+                    else:
+                        skipped_datasets += 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. {skipped_datasets} 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"
+                )
+
+            ingestion_job_summary.increase_skipped_datasets(skipped_datasets)
+
+            if ingestion_job_summary.task_count() >= MAX_TASKS_PER_CHUNK:
+                finish_task_timer()
+                ingestion_job_summary.set_finished()
+                yield ingestion_job_summary
+
+                # Start a new one
+                is_first_chunk = False
+                ingestion_job_summary = IngestionJobSummary.new(ingestion_job=self)
+
+                # We will resume tasks, start timer right away
+                finish_task_timer = ingestion_job_summary.start_timing("tasks")
+
+        if ingestion_job_summary.task_count() > 0 or is_first_chunk:
+            # When there is interesting information to store, or there was no data at all, store it
+            finish_task_timer()
+            ingestion_job_summary.set_finished()
+            yield ingestion_job_summary

ingestify/domain/models/ingestion/ingestion_job_summary.py

@@ -0,0 +1,123 @@
+import uuid
+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():.2f}sec"
+
+
+class IngestionJobSummary(BaseModel):
+    ingestion_job_summary_id: str
+    ingestion_job_id: str
+
+    # From the IngestionPlan
+    provider: str
+    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)
+
+    skipped_datasets: int = 0
+    failed_tasks: int = 0
+    successful_tasks: int = 0
+    ignored_successful_tasks: int = 0
+
+    @classmethod
+    def new(cls, ingestion_job: "IngestionJob"):
+        args = dict(
+            ingestion_job_summary_id=str(uuid.uuid1()),
+            ingestion_job_id=ingestion_job.ingestion_job_id,
+            provider=ingestion_job.ingestion_plan.source.provider,
+            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 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)
+
+    def increase_skipped_datasets(self, skipped_datasets: int):
+        self.skipped_datasets += skipped_datasets
+
+    def task_count(self):
+        return len(self.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"        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"  - 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.successful_tasks}")
+        print(f"    - Skipped datasets: {self.skipped_datasets}")
+        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