ingestify 0.6.4__py3-none-any.whl → 0.7.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.6.4"
+__version__ = "0.7.0"

ingestify/application/dataset_store.py

@@ -1,13 +1,22 @@
 import gzip
-import hashlib
 import logging
-import mimetypes
 import os
 import shutil
-from dataclasses import asdict
+from contextlib import contextmanager
+import threading
 from io import BytesIO
 
-from typing import Dict, List, Optional, Union, Callable, BinaryIO, Awaitable
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Union,
+    Callable,
+    BinaryIO,
+    Awaitable,
+    NewType,
+    Iterable,
+)
 
 from ingestify.domain.models.dataset.dataset import DatasetState
 from ingestify.domain.models.dataset.events import RevisionAdded, MetadataUpdated
@@ -19,7 +28,6 @@ from ingestify.domain.models import (
     Dataset,
     DatasetCollection,
     DatasetRepository,
-    DatasetResource,
     DraftFile,
     File,
     LoadedFile,
@@ -34,6 +42,67 @@ from ingestify.utils import utcnow
 
 logger = logging.getLogger(__name__)
 
+# Type definition for dataset state parameters that can be strings or DatasetState objects
+DatasetStateParam = NewType(
+    "DatasetStateParam", Union[str, Iterable[str], DatasetState, Iterable[DatasetState]]
+)
+
+
+def normalize_dataset_state(
+    dataset_state: Optional[DatasetStateParam],
+) -> Optional[List[DatasetState]]:
+    """
+    Normalize dataset_state parameter to a list of DatasetState objects.
+
+    Args:
+        dataset_state: Can be None, a string, a DatasetState enum,
+                      or a list of strings or DatasetState enums
+
+    Returns:
+        None if input is None, otherwise a list of DatasetState objects
+
+    Raises:
+        ValueError: If an invalid state value is provided
+        TypeError: If dataset_state contains elements of invalid types
+        Warning: If an empty list is provided
+    """
+    if dataset_state is None:
+        return None
+
+    # Check for empty list
+    if isinstance(dataset_state, list) and len(dataset_state) == 0:
+        logger.warning(
+            "Empty list provided for dataset_state, this will not filter any states"
+        )
+        return None
+
+    normalized_states = []
+    states_to_process = (
+        [dataset_state] if not isinstance(dataset_state, list) else dataset_state
+    )
+
+    for state in states_to_process:
+        if isinstance(state, str):
+            # Handle case-insensitive string matching
+            try:
+                # Try to match the string to a DatasetState enum value
+                normalized_state = DatasetState(state.upper())
+                normalized_states.append(normalized_state)
+            except ValueError:
+                valid_states = ", ".join([s.value for s in DatasetState])
+                raise ValueError(
+                    f"Invalid dataset state: '{state}'. Valid states are: {valid_states}"
+                )
+        elif isinstance(state, DatasetState):
+            # Already a DatasetState enum, just add it
+            normalized_states.append(state)
+        else:
+            raise TypeError(
+                f"Dataset state must be a string or DatasetState enum, got {type(state).__name__}"
+            )
+
+    return normalized_states
+
 
 class DatasetStore:
     def __init__(
@@ -48,8 +117,26 @@ class DatasetStore:
         self.bucket = bucket
         self.event_bus: Optional[EventBus] = None
 
-    # def __getstate__(self):
-    #     return {"file_repository": self.file_repository, "bucket": self.bucket}
+        # Pass current version to repository for validation/migration
+        from ingestify import __version__
+
+        self.dataset_repository.ensure_compatible_version(__version__)
+
+    @property
+    def _thread_local(self):
+        if not hasattr(self, "_thread_local_"):
+            self._thread_local_ = threading.local()
+        return self._thread_local_
+
+    def __getstate__(self):
+        """When pickling this instance, don't pass EventBus. EventBus can contain all
+        kind of dispatchers, which may, or may not can be pickled."""
+        return {
+            "dataset_repository": self.dataset_repository,
+            "storage_compression_method": self.storage_compression_method,
+            "file_repository": self.file_repository,
+            "bucket": self.bucket,
+        }
 
     def set_event_bus(self, event_bus: EventBus):
         self.event_bus = event_bus
@@ -58,6 +145,34 @@ class DatasetStore:
         if self.event_bus:
             self.event_bus.dispatch(event)
 
+    @contextmanager
+    def with_file_cache(self):
+        """Context manager to enable file caching during its scope.
+
+        Files loaded within this context will be cached and reused,
+        avoiding multiple downloads of the same file.
+
+        Example:
+            # Without caching (loads files twice)
+            analyzer1 = StatsAnalyzer(store, dataset)
+            analyzer2 = VisualizationTool(store, dataset)
+
+            # With caching (files are loaded once and shared)
+            with store.with_file_cache():
+                analyzer1 = StatsAnalyzer(store, dataset)
+                analyzer2 = VisualizationTool(store, dataset)
+        """
+        # Enable caching for this thread
+        self._thread_local.use_file_cache = True
+        self._thread_local.file_cache = {}
+
+        try:
+            yield
+        finally:
+            # Disable caching for this thread
+            self._thread_local.use_file_cache = False
+            self._thread_local.file_cache = {}
+
     def save_ingestion_job_summary(self, ingestion_job_summary):
         self.dataset_repository.save_ingestion_job_summary(ingestion_job_summary)
 
@@ -67,6 +182,9 @@ class DatasetStore:
         provider: Optional[str] = None,
         dataset_id: Optional[str] = None,
         metadata_only: Optional[bool] = False,
+        page: Optional[int] = None,
+        page_size: Optional[int] = None,
+        dataset_state: Optional[DatasetStateParam] = None,
         **selector,
     ) -> DatasetCollection:
         if "selector" in selector:
@@ -82,6 +200,9 @@ class DatasetStore:
                 # Convert all selector dicts to Selectors
                 selector = [Selector(_) for _ in selector]
 
+        # Normalize dataset_state to a list of DatasetState objects
+        normalized_dataset_state = normalize_dataset_state(dataset_state)
+
         dataset_collection = self.dataset_repository.get_dataset_collection(
             bucket=self.bucket,
             dataset_type=dataset_type,
@@ -89,9 +210,89 @@ class DatasetStore:
             provider=provider,
             metadata_only=metadata_only,
             selector=selector,
+            dataset_state=normalized_dataset_state,
+            page=page,
+            page_size=page_size,
         )
         return dataset_collection
 
+    def iter_dataset_collection_batches(
+        self,
+        dataset_type: Optional[str] = None,
+        provider: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        batch_size: int = 1000,
+        yield_dataset_collection: bool = False,
+        dataset_state: Optional[DatasetStateParam] = None,
+        **selector,
+    ):
+        """
+        Iterate through all datasets matching the criteria with automatic pagination.
+
+        Examples:
+        ```
+        # Iterate through individual datasets
+        for dataset in store.iter_dataset_collection_batches(dataset_type="match", provider="statsbomb"):
+            process(dataset)
+
+        # Iterate through DatasetCollection objects (pages)
+        for collection in store.iter_dataset_collection(
+            dataset_type="match",
+            provider="statsbomb",
+            yield_dataset_collection=True
+        ):
+            process_collection(collection)
+
+        # Filter by dataset state
+        for dataset in store.iter_dataset_collection(
+            dataset_type="match",
+            dataset_state="COMPLETE"  # Can also use DatasetState.COMPLETE or ["COMPLETE", "PARTIAL"]
+        ):
+            process_completed_dataset(dataset)
+        ```
+
+        Args:
+            dataset_type: Optional dataset type filter
+            provider: Optional provider filter
+            dataset_id: Optional dataset ID filter
+            batch_size: Number of datasets to fetch per batch
+            yield_dataset_collection: If True, yields entire DatasetCollection objects
+                                     instead of individual Dataset objects
+            dataset_state: Optional filter for dataset state. Can be a string, DatasetState enum,
+                          or a list of strings or DatasetState enums
+            **selector: Additional selector criteria
+
+        Yields:
+            If yield_dataset_collection is False (default): Dataset objects one by one
+            If yield_dataset_collection is True: DatasetCollection objects (pages)
+        """
+        page = 1
+        while True:
+            collection = self.get_dataset_collection(
+                dataset_type=dataset_type,
+                provider=provider,
+                dataset_id=dataset_id,
+                page=page,
+                page_size=batch_size,
+                dataset_state=dataset_state,
+                **selector,
+            )
+
+            if not collection or len(collection) == 0:
+                break
+
+            if yield_dataset_collection:
+                yield collection
+            else:
+                for dataset in collection:
+                    yield dataset
+
+            # If we got fewer results than page_size, we've reached the end
+            if len(collection) < batch_size:
+                break
+
+            page += 1
+
     #
     # def destroy_dataset(self, dataset_id: str):
     #     dataset = self.dataset_repository.
@@ -311,10 +512,21 @@ class DatasetStore:
                     self.file_repository.load_content(storage_path=file_.storage_path)
                 )
 
-            loaded_file = LoadedFile(
-                stream_=get_stream if lazy else get_stream(file),
-                **file.model_dump(),
-            )
+            def make_loaded_file():
+                return LoadedFile(
+                    stream_=get_stream if lazy else get_stream(file),
+                    **file.model_dump(),
+                )
+
+            # Using getattr with a default value of False - simple one-liner
+            if getattr(self._thread_local, "use_file_cache", False):
+                key = (dataset.dataset_id, current_revision.revision_id, file.file_id)
+                if key not in self._thread_local.file_cache:
+                    self._thread_local.file_cache[key] = make_loaded_file()
+                loaded_file = self._thread_local.file_cache[key]
+            else:
+                loaded_file = make_loaded_file()
+
             files[file.file_id] = loaded_file
         return FileCollection(files, auto_rewind=auto_rewind)
 
@@ -324,9 +536,14 @@ class DatasetStore:
             from kloppy import statsbomb
 
             try:
+                three_sixty_data = None
+                if tmp_file := files.get_file("360-frames"):
+                    three_sixty_data = tmp_file.stream
+
                 return statsbomb.load(
                     event_data=(files.get_file("events")).stream,
                     lineup_data=(files.get_file("lineups")).stream,
+                    three_sixty_data=three_sixty_data,
                     **kwargs,
                 )
             except Exception as e:

ingestify/application/ingestion_engine.py

@@ -1,15 +1,31 @@
 import itertools
 import logging
-from typing import Optional, List
-
+import threading
+from queue import Queue
+from typing import Optional, List, Union, Any, Iterator, TypedDict
 
 from .loader import Loader
 from .dataset_store import DatasetStore
 from ingestify.domain.models.ingestion.ingestion_plan import IngestionPlan
+from ingestify.domain.models import Dataset
+from ..domain.models.dataset.events import (
+    DatasetSkipped,
+    RevisionAdded,
+    SelectorSkipped,
+)
+from ..domain.models.event import DomainEvent
 
 logger = logging.getLogger(__name__)
 
 
+class AutoIngestConfig(TypedDict, total=False):
+    """Configuration options for auto-ingestion feature."""
+
+    enabled: bool
+    streaming: bool
+    use_open_data: bool
+
+
 class IngestionEngine:
     def __init__(self, store: DatasetStore):
 
@@ -26,8 +42,73 @@ class IngestionEngine:
         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)
+        dataset_type: Optional[str] = None,
+        auto_ingest_config: Optional[AutoIngestConfig] = None,
+        async_yield_events: bool = False,
+        **selector_filters,
+    ) -> Optional[Iterator[DomainEvent]]:
+        """
+        Execute data ingestion from configured sources.
+
+        Args:
+            dry_run: If True, perform validation but skip actual data ingestion
+            provider: Filter ingestion to specific data provider
+            source: Filter ingestion to specific source name
+            dataset_type: Filter ingestion to specific dataset type (e.g., 'match', 'lineups')
+            auto_ingest_config: Configuration for auto-discovery of ingestion plans
+            async_yield_events: If True, run ingestion in background and yield domain events
+            **selector_filters: Additional selector criteria (e.g., competition_id=43)
+
+        Returns:
+            Iterator of DomainEvent objects if async_yield_events=True, None otherwise
+
+        Examples:
+            # Standard synchronous ingestion
+            engine.load(dataset_type="match", competition_id=43)
+
+            # Real-time event streaming during ingestion
+            for event in engine.load(async_yield_events=True):
+                if isinstance(event, DatasetSkipped):
+                    print(f"Dataset ready: {event.dataset.name}")
+        """
+
+        def do_load():
+            self.loader.collect_and_run(
+                dry_run=dry_run,
+                provider=provider,
+                source=source,
+                dataset_type=dataset_type,
+                auto_ingest_config=auto_ingest_config or {},
+                **selector_filters,
+            )
+
+        if async_yield_events:
+            queue = Queue()
+
+            def load_in_background():
+                unregister = self.store.event_bus.register_queue(queue)
+                do_load()
+                unregister()
+
+                # Done.
+                queue.put(None)
+
+            thread = threading.Thread(target=load_in_background)
+            thread.start()
+
+            def _iter():
+                # Required to prevent this function to be a generator. Otherwise,
+                # do_load won't be called when async_yield_events=False
+                while True:
+                    event = queue.get()
+                    if event is None:
+                        break
+
+                    yield event
+
+            return _iter()
+        else:
+            do_load()
 
     def list_datasets(self, as_count: bool = False):
         """Consider moving this to DataStore"""
@@ -52,16 +133,157 @@ class IngestionEngine:
                     print(f"  {dataset_type}:")
                     for dataset in datasets_per_type:
                         print(
-                            f"    {dataset.identifier}: {dataset.name} / {dataset.state}   {dataset.dataset_id}"
+                            f"    {dataset.identifier}\t{dataset.dataset_id}\t{dataset.name} / {dataset.state}"
                         )
             # print(dataset.dataset_id)
 
     def destroy_dataset(
         self, dataset_id: Optional[str] = None, **selector
     ) -> List[str]:
-        datasets = self.store.get_dataset_collection(dataset_id=dataset_id, **selector)
+        dataset_collection = self.store.iter_dataset_collection_batches(
+            dataset_id=dataset_id,
+            **selector,
+        )
         dataset_ids = []
-        for dataset in datasets:
+        for dataset in dataset_collection:
             self.store.destroy_dataset(dataset)
             dataset_ids.append(dataset.dataset_id)
         return dataset_ids
+
+    def iter_datasets(
+        self,
+        auto_ingest: Union[bool, AutoIngestConfig] = False,
+        dataset_type: Optional[str] = None,
+        provider: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        batch_size: int = 1000,
+        yield_dataset_collection: bool = False,
+        dataset_state: Optional[Any] = None,
+        **selector_filters,
+    ) -> Iterator[Dataset]:
+        """
+        Iterate over datasets with optional auto-ingestion.
+
+        This method combines dataset discovery, ingestion, and retrieval in a single
+        streaming interface. When auto_ingest=True, it will first run the ingestion
+        pipeline to discover and ingest matching datasets before yielding results.
+
+        Examples:
+            # Basic iteration over existing datasets
+            for dataset in engine.iter_datasets(provider="statsbomb"):
+                process(dataset)
+
+            # Auto-ingest new data matching criteria before iteration
+            for dataset in engine.iter_datasets(
+                auto_ingest=True,
+                provider="statsbomb",
+                competition_id=11
+            ):
+                process(dataset)  # Includes newly ingested datasets
+
+            # Real-time streaming with open data auto-discovery
+            for dataset in engine.iter_datasets(
+                auto_ingest={
+                    "streaming": True,
+                    "use_open_data": True
+                },
+                dataset_type="match"
+            ):
+                process(dataset)  # Yields datasets as they become available
+
+        Args:
+            auto_ingest: Enable auto-ingestion before yielding datasets.
+                        Can be True/False or AutoIngestConfig dict with options:
+                        - enabled: bool (default: True)
+                        - streaming: bool (default: False) - enables real-time event streaming
+                        - use_open_data: bool (default: False) - auto-discover open data sources
+            dataset_type: Filter by dataset type (e.g., "match", "competition")
+            provider: Filter by data provider (e.g., "statsbomb", "wyscout")
+            dataset_id: Filter by specific dataset ID
+            batch_size: Number of datasets to fetch per batch for pagination
+            yield_dataset_collection: If True, yield DatasetCollection objects instead of individual datasets
+            dataset_state: Filter by dataset state (e.g., "COMPLETE", "PARTIAL")
+            **selector_filters: Additional selector criteria (competition_id, season_id, match_id, etc.)
+
+        Yields:
+            Dataset objects matching the specified criteria. If auto_ingest=True,
+            includes both existing datasets and newly ingested ones.
+
+        Note:
+            Auto-ingestion will only discover datasets that match configured
+            IngestionPlans. Requests outside the scope of existing plans will
+            not trigger ingestion.
+        """
+        # Parse auto_ingest config
+        if isinstance(auto_ingest, dict):
+            auto_ingest_enabled = auto_ingest.get("enabled", True)
+            auto_ingest_config = auto_ingest
+        else:
+            auto_ingest_enabled = bool(auto_ingest)
+            auto_ingest_config = {}
+
+        # Run auto-ingestion if enabled
+        if auto_ingest_enabled:
+            if auto_ingest_config.get("streaming", False):
+                if yield_dataset_collection:
+                    raise ValueError(
+                        "Cannot yield_dataset_collection when "
+                        "auto_ingest_enabled. In case of streaming mode"
+                    )
+
+                # Start background loading immediately - don't return a generator
+                event_iter = self.load(
+                    provider=provider,
+                    dataset_type=dataset_type,
+                    auto_ingest_config=auto_ingest_config,
+                    async_yield_events=True,
+                    **selector_filters,
+                )
+
+                for event in event_iter:
+                    if isinstance(event, DatasetSkipped):
+                        yield event.dataset
+                    elif isinstance(event, RevisionAdded):
+                        yield event.dataset
+                    elif isinstance(event, SelectorSkipped):
+                        yield from self.store.iter_dataset_collection_batches(
+                            dataset_type=dataset_type,
+                            provider=provider,
+                            batch_size=batch_size,
+                            # We can't yield Dataset (from DatasetSkipped and RevisionAdded) and
+                            # DatasetCollection in the same run
+                            yield_dataset_collection=False,
+                            dataset_state=dataset_state,
+                            **event.selector.filtered_attributes,
+                        )
+                return
+            else:
+                self.load(
+                    provider=provider,
+                    dataset_type=dataset_type,
+                    auto_ingest_config=auto_ingest_config,
+                    **selector_filters,
+                )
+
+        yield from self.store.iter_dataset_collection_batches(
+            dataset_type=dataset_type,
+            provider=provider,
+            dataset_id=dataset_id,
+            batch_size=batch_size,
+            yield_dataset_collection=yield_dataset_collection,
+            dataset_state=dataset_state,
+            **selector_filters,
+        )
+
+    def load_dataset_with_kloppy(self, dataset: Dataset, **kwargs):
+        """
+        Load a dataset using kloppy.
+
+        Args:
+            dataset: The dataset to load
+            **kwargs: Additional arguments passed to kloppy's load function
+
+        Returns:
+            Kloppy dataset object
+        """
+        return self.store.load_with_kloppy(dataset, **kwargs)