ingestify 0.1.3__py3-none-any.whl → 0.3.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.1.3"
+__version__ = "0.3.0"

ingestify/application/dataset_store.py

@@ -5,13 +5,14 @@ import mimetypes
 import os
 import shutil
 from dataclasses import asdict
-from io import BytesIO, StringIO
+from io import BytesIO
 
-from typing import Dict, List, Optional, Union, Callable, BinaryIO
+from typing import Dict, List, Optional, Union, Callable, BinaryIO, Awaitable
 
 from ingestify.domain.models.dataset.dataset import DatasetState
 from ingestify.domain.models.dataset.events import RevisionAdded, MetadataUpdated
 from ingestify.domain.models.dataset.file_collection import FileCollection
+from ingestify.domain.models.dataset.revision import RevisionSource
 from ingestify.domain.models.event import EventBus
 from ingestify.domain.models import (
     Dataset,
@@ -27,7 +28,7 @@ from ingestify.domain.models import (
     Revision,
     DatasetCreated,
 )
-from ingestify.utils import utcnow, map_in_pool
+from ingestify.utils import utcnow
 
 
 logger = logging.getLogger(__name__)
@@ -56,11 +57,16 @@ class DatasetStore:
         if self.event_bus:
             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()
+
     def get_dataset_collection(
         self,
         dataset_type: Optional[str] = None,
         provider: Optional[str] = None,
         dataset_id: Optional[str] = None,
+        metadata_only: Optional[bool] = False,
         **selector,
     ) -> DatasetCollection:
         if "selector" in selector:
@@ -81,6 +87,7 @@ class DatasetStore:
             dataset_type=dataset_type,
             dataset_id=dataset_id,
             provider=provider,
+            metadata_only=metadata_only,
             selector=selector,
         )
         return dataset_collection
@@ -107,7 +114,9 @@ class DatasetStore:
 
         return stream, storage_size, suffix
 
-    def _prepare_read_stream(self) -> tuple[Callable[[BinaryIO], BytesIO], str]:
+    def _prepare_read_stream(
+        self,
+    ) -> tuple[Callable[[BinaryIO], Awaitable[BytesIO]], str]:
         if self.storage_compression_method == "gzip":
 
             def reader(fh: BinaryIO) -> BytesIO:
@@ -168,7 +177,11 @@ class DatasetStore:
         return modified_files_
 
     def add_revision(
-        self, dataset: Dataset, files: Dict[str, DraftFile], description: str = "Update"
+        self,
+        dataset: Dataset,
+        files: Dict[str, DraftFile],
+        revision_source: RevisionSource,
+        description: str = "Update",
     ):
         """
         Create new revision first, so FileRepository can use
@@ -182,46 +195,53 @@ class DatasetStore:
             # It can happen an API tells us data is changed, but it was not changed. In this case
             # we decide to ignore it.
             # Make sure there are files changed before creating a new revision
-            dataset.add_revision(
-                Revision(
-                    revision_id=revision_id,
-                    created_at=created_at,
-                    description=description,
-                    modified_files=persisted_files_,
-                )
+            revision = Revision(
+                revision_id=revision_id,
+                created_at=created_at,
+                description=description,
+                modified_files=persisted_files_,
+                source=revision_source,
             )
 
+            dataset.add_revision(revision)
+
             self.dataset_repository.save(bucket=self.bucket, dataset=dataset)
             self.dispatch(RevisionAdded(dataset=dataset))
             logger.info(
                 f"Added a new revision to {dataset.identifier} -> {', '.join([file.file_id for file in persisted_files_])}"
             )
-            return True
         else:
             logger.info(
                 f"Ignoring a new revision without changed files -> {dataset.identifier}"
             )
-            return False
+            revision = None
+
+        return revision
 
     def update_dataset(
         self,
         dataset: Dataset,
-        dataset_resource: DatasetResource,
+        name: str,
+        state: DatasetState,
+        metadata: dict,
         files: Dict[str, DraftFile],
+        revision_source: RevisionSource,
     ):
         """The add_revision will also save the dataset."""
         metadata_changed = False
-        if dataset.update_from_resource(dataset_resource):
+        if dataset.update_metadata(name, metadata, state):
             self.dataset_repository.save(bucket=self.bucket, dataset=dataset)
             metadata_changed = True
 
-        self.add_revision(dataset, files)
+        revision = self.add_revision(dataset, files, revision_source)
 
         if metadata_changed:
             # Dispatch after revision added. Otherwise, the downstream handlers are not able to see
             # the new revision
             self.dispatch(MetadataUpdated(dataset=dataset))
 
+        return revision
+
     def destroy_dataset(self, dataset: Dataset):
         # TODO: remove files. Now we leave some orphaned files around
         self.dataset_repository.destroy(dataset)
@@ -235,6 +255,7 @@ class DatasetStore:
         state: DatasetState,
         metadata: dict,
         files: Dict[str, DraftFile],
+        revision_source: RevisionSource,
         description: str = "Create",
     ):
         now = utcnow()
@@ -251,9 +272,10 @@ class DatasetStore:
             created_at=now,
             updated_at=now,
         )
-        self.add_revision(dataset, files, description)
+        revision = self.add_revision(dataset, files, revision_source, description)
 
         self.dispatch(DatasetCreated(dataset=dataset))
+        return revision
 
     def load_files(
         self,
@@ -271,20 +293,9 @@ class DatasetStore:
                 continue
 
             def get_stream(file_):
-                revision_id = file_.revision_id
-                if revision_id is None:
-                    revision_id = current_revision.revision_id
-
                 return reader(
                     self.file_repository.load_content(
-                        bucket=self.bucket,
-                        dataset=dataset,
-                        # When file.revision_id is set we must use it.
-                        revision_id=revision_id,
-                        filename=file_.file_id
-                        + "."
-                        + file_.data_serialization_format
-                        + suffix,
+                        bucket=self.bucket, storage_path=file_.storage_path
                     )
                 )
 
@@ -302,8 +313,8 @@ class DatasetStore:
 
             try:
                 return statsbomb.load(
-                    event_data=files.get_file("events").stream,
-                    lineup_data=files.get_file("lineups").stream,
+                    event_data=(files.get_file("events")).stream,
+                    lineup_data=(files.get_file("lineups")).stream,
                     **kwargs,
                 )
             except Exception as e:
@@ -333,7 +344,7 @@ class DatasetStore:
     #         filename=filename,
     #     )
 
-    def map(
-        self, fn, dataset_collection: DatasetCollection, processes: Optional[int] = None
-    ):
-        return map_in_pool(fn, dataset_collection, processes)
+    # def map(
+    #     self, fn, dataset_collection: DatasetCollection, processes: Optional[int] = None
+    # ):
+    #     return map_in_pool(fn, dataset_collection, processes)

ingestify/application/ingestion_engine.py

@@ -5,7 +5,7 @@ from typing import Optional, List
 
 from .loader import Loader
 from .dataset_store import DatasetStore
-from ..domain.models.extract_job import ExtractJob
+from ingestify.domain.models.ingestion.ingestion_plan import IngestionPlan
 
 logger = logging.getLogger(__name__)
 
@@ -18,8 +18,8 @@ class IngestionEngine:
         self.store = store
         self.loader = Loader(self.store)
 
-    def add_extract_job(self, extract_job: ExtractJob):
-        self.loader.add_extract_job(extract_job)
+    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)

ingestify/application/loader.py

@@ -1,19 +1,15 @@
-import itertools
-import json
 import logging
 import platform
-from multiprocessing import set_start_method, cpu_count
+import uuid
+from multiprocessing import set_start_method
 from typing import List, Optional
 
-from ingestify.domain.models import Dataset, Identifier, Selector, Source, Task, TaskSet
-from ingestify.utils import map_in_pool, TaskExecutor, chunker
+from ingestify.domain.models import Selector
+from ingestify.utils import TaskExecutor
 
 from .dataset_store import DatasetStore
-from .. import DatasetResource, retrieve_http
-from ..domain import DraftFile
-from ..domain.models.data_spec_version_collection import DataSpecVersionCollection
-from ..domain.models.extract_job import ExtractJob
-from ..domain.models.resources.dataset_resource import FileResource
+from ingestify.domain.models.ingestion.ingestion_plan import IngestionPlan
+from ..domain.models.ingestion.ingestion_job import IngestionJob
 from ..exceptions import ConfigurationError
 
 if platform.system() == "Darwin":
@@ -25,176 +21,54 @@ else:
 logger = logging.getLogger(__name__)
 
 
-DEFAULT_CHUNK_SIZE = 1000
-
-
-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",
-            **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
-
-    def run(self):
-        self.store.update_dataset(
-            dataset=self.dataset,
-            dataset_resource=self.dataset_resource,
-            files={
-                file_id: load_file(file_resource, dataset=self.dataset)
-                for file_id, file_resource in self.dataset_resource.files.items()
-            },
-        )
-
-    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
-
-    def run(self):
-        self.store.create_dataset(
-            dataset_type=self.dataset_resource.dataset_type,
-            provider=self.dataset_resource.provider,
-            dataset_identifier=Identifier(**self.dataset_resource.dataset_resource_id),
-            name=self.dataset_resource.name,
-            state=self.dataset_resource.state,
-            metadata=self.dataset_resource.metadata,
-            files={
-                file_id: load_file(file_resource)
-                for file_id, file_resource in self.dataset_resource.files.items()
-            },
-        )
-
-    def __repr__(self):
-        return f"CreateDatasetTask({self.dataset_resource.provider} -> {self.dataset_resource.dataset_resource_id})"
-
-
 class Loader:
     def __init__(self, store: DatasetStore):
         self.store = store
-        self.extract_jobs: List[ExtractJob] = []
+        self.ingestion_plans: List[IngestionPlan] = []
 
-    def add_extract_job(self, extract_job: ExtractJob):
-        self.extract_jobs.append(extract_job)
+    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):
-        total_dataset_count = 0
-
         # First collect all selectors, before discovering datasets
         selectors = {}
-        for extract_job in self.extract_jobs:
+        for ingestion_plan in self.ingestion_plans:
+            logger.info(f"Determining selectors for {ingestion_plan}")
+
             if provider is not None:
-                if extract_job.source.provider != provider:
+                if ingestion_plan.source.provider != provider:
                     logger.info(
-                        f"Skipping {extract_job } because provider doesn't match '{provider}'"
+                        f"Skipping {ingestion_plan} because provider doesn't match '{provider}'"
                     )
                     continue
 
             static_selectors = [
                 selector
-                for selector in extract_job.selectors
+                for selector in ingestion_plan.selectors
                 if not selector.is_dynamic
             ]
             dynamic_selectors = [
-                selector for selector in extract_job.selectors if selector.is_dynamic
+                selector for selector in ingestion_plan.selectors if selector.is_dynamic
             ]
 
             no_selectors = len(static_selectors) == 1 and not bool(static_selectors[0])
             if dynamic_selectors or no_selectors:
-                if hasattr(extract_job.source, "discover_selectors"):
+                if hasattr(ingestion_plan.source, "discover_selectors"):
                     logger.debug(
-                        f"Discovering selectors from {extract_job.source.__class__.__name__}"
+                        f"Discovering selectors from {ingestion_plan.source.__class__.__name__}"
                     )
 
                     # TODO: consider making this lazy and fetch once per Source instead of
-                    #       once per ExtractJob
-                    all_selectors = extract_job.source.discover_selectors(
-                        extract_job.dataset_type
+                    #       once per IngestionPlan
+                    all_selectors = ingestion_plan.source.discover_selectors(
+                        ingestion_plan.dataset_type
                     )
                     if no_selectors:
                         # When there were no selectors specified, just use all of them
                         extra_static_selectors = [
                             Selector.build(
                                 job_selector,
-                                data_spec_versions=extract_job.data_spec_versions,
+                                data_spec_versions=ingestion_plan.data_spec_versions,
                             )
                             for job_selector in all_selectors
                         ]
@@ -205,7 +79,7 @@ class Loader:
                             dynamic_job_selectors = [
                                 Selector.build(
                                     job_selector,
-                                    data_spec_versions=extract_job.data_spec_versions,
+                                    data_spec_versions=ingestion_plan.data_spec_versions,
                                 )
                                 for job_selector in all_selectors
                                 if dynamic_selector.is_match(job_selector)
@@ -216,7 +90,7 @@ class Loader:
                     static_selectors.extend(extra_static_selectors)
 
                     logger.info(
-                        f"Discovered {len(extra_static_selectors)} selectors from {extract_job.source.__class__.__name__}"
+                        f"Discovered {len(extra_static_selectors)} selectors from {ingestion_plan.source.__class__.__name__}"
                     )
                 else:
                     if not no_selectors:
@@ -224,112 +98,68 @@ class Loader:
                         # later on
                         raise ConfigurationError(
                             f"Dynamic selectors cannot be used for "
-                            f"{extract_job.source.__class__.__name__} because it doesn't support"
+                            f"{ingestion_plan.source.__class__.__name__} because it doesn't support"
                             f" selector discovery"
                         )
 
             # Merge selectors when source, dataset_type and actual selector is the same. This makes
             # sure there will be only 1 dataset for this combination
             for selector in static_selectors:
-                key = (extract_job.source.name, extract_job.dataset_type, selector.key)
+                key = (
+                    ingestion_plan.source.name,
+                    ingestion_plan.dataset_type,
+                    selector.key,
+                )
                 if existing_selector := selectors.get(key):
                     existing_selector[1].data_spec_versions.merge(
                         selector.data_spec_versions
                     )
                 else:
-                    selectors[key] = (extract_job, selector)
-
-        def run_task(task):
-            logger.info(f"Running task {task}")
-            task.run()
-
-        for extract_job, selector in selectors.values():
-            logger.debug(
-                f"Discovering datasets from {extract_job.source.__class__.__name__} using selector {selector}"
+                    selectors[key] = (ingestion_plan, selector)
+
+        """
+            Data is denormalized:
+            
+            It actually looks like:
+                - IngestionPlan #1
+                    - Selector 1.1
+                    - Selector 1.2
+                    - Selector 1.3
+                - IngestionPlan #2
+                    - Selector 2.1
+                    - Selector 2.2
+                    
+            We process this as:
+            - IngestionPlan #1, Selector 1.1
+            - IngestionPlan #1, Selector 1.2
+            - IngestionPlan #1, Selector 1.3
+            - IngestionPlan #2, Selector 2.1
+            - IngestionPlan #2, Selector 2.2 
+            
+            IngestionJobSummary holds the summary for an IngestionPlan and a single Selector
+        """
+        for ingestion_plan, selector in selectors.values():
+            logger.info(
+                f"Discovering datasets from {ingestion_plan.source.__class__.__name__} using selector {selector}"
             )
 
-            dataset_collection_metadata = self.store.get_dataset_collection(
-                dataset_type=extract_job.dataset_type,
-                data_spec_versions=selector.data_spec_versions,
+            ingestion_job = IngestionJob(
+                ingestion_job_id=str(uuid.uuid1()),
+                ingestion_plan=ingestion_plan,
                 selector=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
-            datasets = extract_job.source.find_datasets(
-                dataset_type=extract_job.dataset_type,
-                data_spec_versions=selector.data_spec_versions,
-                dataset_collection_metadata=dataset_collection_metadata,
-                **selector.custom_attributes,
             )
 
-            batches = to_batches(datasets)
-
-            for batch in batches:
-                dataset_identifiers = [
-                    Identifier.create_from_selector(
-                        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 = self.store.get_dataset_collection(
-                    dataset_type=extract_job.dataset_type,
-                    # Assume all DatasetResources share the same provider
-                    provider=batch[0].provider,
-                    selector=dataset_identifiers,
-                )
-
-                skip_count = 0
-                total_dataset_count += len(dataset_identifiers)
-
-                task_set = TaskSet()
-                for dataset_resource in batch:
-                    dataset_identifier = Identifier.create_from_selector(
-                        selector, **dataset_resource.dataset_resource_id
-                    )
-
-                    if dataset := dataset_collection.get(dataset_identifier):
-                        if extract_job.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=self.store,
-                                )
-                            )
-                        else:
-                            skip_count += 1
-                    else:
-                        if extract_job.fetch_policy.should_fetch(dataset_resource):
-                            task_set.add(
-                                CreateDatasetTask(
-                                    dataset_resource=dataset_resource,
-                                    store=self.store,
-                                )
-                            )
-                        else:
-                            skip_count += 1
-
-                if task_set:
-                    logger.info(
-                        f"Discovered {len(dataset_identifiers)} datasets from {extract_job.source.__class__.__name__} "
-                        f"using selector {selector} => {len(task_set)} tasks. {skip_count} skipped."
-                    )
-                    logger.info(f"Running {len(task_set)} tasks")
-                    with TaskExecutor(dry_run=dry_run) as task_executor:
-                        task_executor.run(run_task, task_set)
-                else:
-                    logger.info(
-                        f"Discovered {len(dataset_identifiers)} datasets from {extract_job.source.__class__.__name__} "
-                        f"using selector {selector} => nothing to do"
-                    )
+            with TaskExecutor(dry_run=dry_run) as task_executor:
+                for ingestion_job_summary in ingestion_job.execute(
+                    self.store, task_executor=task_executor
+                ):
+                    # TODO: handle task_summaries
+                    #       Summarize to a IngestionJobSummary, and save to a database. This Summary can later be used in a
+                    #       next run to determine where to resume.
+                    # TODO 2: Do we want to add additional information from the summary back to the Task, so it can use
+                    #      extra information to determine how/where to resume
+                    ingestion_job_summary.output_report()
+                    logger.info(f"Storing IngestionJobSummary")
+                    self.store.save_ingestion_job_summary(ingestion_job_summary)
 
         logger.info("Done")

ingestify/domain/models/__init__.py

@@ -11,10 +11,8 @@ from .dataset import (
     LoadedFile,
     Selector,
     Revision,
-    dataset_repository_factory,
-    file_repository_factory,
 )
-from .sink import Sink, sink_factory
+from .sink import Sink
 from .source import Source
 from .task import Task, TaskSet
 from .data_spec_version_collection import DataSpecVersionCollection
@@ -35,11 +33,8 @@ __all__ = [
     "FileRepository",
     "FileCollection",
     "DatasetRepository",
-    "dataset_repository_factory",
-    "file_repository_factory",
     "TaskSet",
     "Task",
     "Sink",
-    "sink_factory",
     "DataSpecVersionCollection",
 ]

ingestify/domain/models/base.py

@@ -0,0 +1,22 @@
+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