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

ingestify/infra/fetch/http.py

@@ -69,7 +69,12 @@ def retrieve_http(
         else:
             raise Exception(f"Don't know how to use {key}")
 
+    ignore_not_found = http_kwargs.pop("ignore_not_found", False)
+
     response = get_session().get(url, headers=headers, **http_kwargs)
+    if response.status_code == 404 and ignore_not_found:
+        return None
+
     response.raise_for_status()
     if response.status_code == 304:
         # Not modified

ingestify/infra/source/statsbomb_github.py

@@ -1,22 +1,22 @@
-import json
 from datetime import datetime
 
 import requests
 
-from ingestify import Source, retrieve_http
-from ingestify.domain import DraftFile
+from ingestify import Source, DatasetResource
 from ingestify.domain.models.dataset.dataset import DatasetState
 
 BASE_URL = "https://raw.githubusercontent.com/statsbomb/open-data/master/data"
+DATA_SPEC_VERSION = "v1-open-data"
 
 
 class StatsbombGithub(Source):
     provider = "statsbomb"
 
-    def discover_selectors(self, dataset_type: str, data_spec_versions: None = None):
+    def discover_selectors(self, dataset_type: str):
         assert dataset_type == "match"
 
         competitions = requests.get(f"{BASE_URL}/competitions.json").json()
+
         return [
             dict(
                 competition_id=competition["competition_id"],
@@ -25,68 +25,81 @@ class StatsbombGithub(Source):
             for competition in competitions
         ]
 
-    def discover_datasets(
+    def find_datasets(
         self,
-        dataset_type,
-        competition_id: str = None,
-        season_id: str = None,
+        dataset_type: str,
+        competition_id: str,
+        season_id: str,
+        match_id: str = None,
         data_spec_versions=None,
+        dataset_collection_metadata=None,
     ):
         assert dataset_type == "match"
 
-        datasets = []
-
         matches = requests.get(
             f"{BASE_URL}/matches/{competition_id}/{season_id}.json"
         ).json()
 
         for match in matches:
-            last_updated = match["last_updated"]
-            if "Z" not in last_updated:
-                # Assume UTC
-                last_updated += "Z"
-
-            last_modified = datetime.fromisoformat(last_updated.replace("Z", "+00:00"))
-
-            dataset = dict(
-                competition_id=competition_id,
-                season_id=season_id,
-                match_id=match["match_id"],
-                _last_modified=last_modified,
-                _match=match,
-                _metadata=match,
-                _state=DatasetState.COMPLETE,
-            )
-            datasets.append(dataset)
-        return datasets
+            if match_id:
+                if match["match_id"] != match_id:
+                    continue
 
-    def fetch_dataset_files(
-        self, dataset_type, identifier, current_revision, data_spec_versions
-    ):
-        assert dataset_type == "match"
+            last_modified = datetime.fromisoformat(match["last_updated"] + "+00:00")
 
-        current_files = current_revision.modified_files_map if current_revision else {}
-        files = {}
-        for filename, url in [
-            ("lineups.json", f"{BASE_URL}/lineups/{identifier.match_id}.json"),
-            ("events.json", f"{BASE_URL}/events/{identifier.match_id}.json"),
-        ]:
-            data_feed_key = filename.split(".")[0]
-            file_id = data_feed_key + "__v1"
-            files[file_id] = retrieve_http(
-                url,
-                current_files.get(filename),
-                file_data_feed_key=data_feed_key,
-                file_data_spec_version="v1",
-                file_data_serialization_format="json",
+            # Open data is always complete.. I guess?
+            state = DatasetState.COMPLETE
+
+            name = (
+                f"{match['match_date']} / "
+                f"{match['home_team']['home_team_name']} - {match['away_team']['away_team_name']}"
+            )
+
+            dataset_resource = DatasetResource(
+                dataset_resource_id=dict(
+                    competition_id=competition_id,
+                    season_id=season_id,
+                    match_id=match["match_id"],
+                ),
+                dataset_type=dataset_type,
+                provider=self.provider,
+                name=name,
+                metadata=match,
+                state=state,
             )
 
-        files["match__v1"] = DraftFile.from_input(
-            json.dumps(identifier._match, indent=4),
-            data_feed_key="match",
-            data_spec_version="v1",
-            data_serialization_format="json",
-            modified_at=None,
-        )
+            dataset_resource.add_file(
+                last_modified=last_modified,
+                data_feed_key="match",
+                data_spec_version=DATA_SPEC_VERSION,
+                json_content=match,
+            )
 
-        return files
+            if state.is_complete:
+                name += f" / {match['home_score']}-{match['away_score']}"
+
+                for data_feed_key in ["lineups", "events"]:
+                    dataset_resource.add_file(
+                        last_modified=last_modified,
+                        data_feed_key=data_feed_key,
+                        data_spec_version=DATA_SPEC_VERSION,
+                        url=f"{BASE_URL}/{data_feed_key}/{match['match_id']}.json",
+                        data_serialization_format="json",
+                    )
+
+                if (
+                    match["last_updated_360"]
+                    and match["match_status_360"] == "available"
+                ):
+                    dataset_resource.add_file(
+                        last_modified=datetime.fromisoformat(
+                            match["last_updated_360"] + "+00:00"
+                        ),
+                        data_feed_key="360-frames",
+                        data_spec_version=DATA_SPEC_VERSION,
+                        url=f"{BASE_URL}/three-sixty/{match['match_id']}.json",
+                        data_serialization_format="json",
+                        http_options={"ignore_not_found": True},
+                    )
+
+            yield dataset_resource

ingestify/infra/store/dataset/__init__.py

@@ -1,2 +0,0 @@
-from .local_dataset_repository import LocalDatasetRepository
-from .sqlalchemy import SqlAlchemyDatasetRepository

ingestify/infra/store/dataset/sqlalchemy/mapping.py

@@ -1,5 +1,7 @@
 import datetime
+from dataclasses import is_dataclass, asdict
 from pathlib import Path
+from typing import Optional
 
 from sqlalchemy import (
     JSON,
@@ -13,11 +15,37 @@ from sqlalchemy import (
     String,
     Table,
     TypeDecorator,
+    Boolean,
 )
 from sqlalchemy.orm import registry, relationship
 
+from ingestify.domain import Selector, Identifier, DataSpecVersionCollection
 from ingestify.domain.models import Dataset, File, Revision
 from ingestify.domain.models.dataset.dataset import DatasetState
+from ingestify.domain.models.ingestion.ingestion_job_summary import (
+    IngestionJobSummary,
+)
+from ingestify.domain.models.task.task_summary import TaskSummary, Operation, TaskStatus
+from ingestify.domain.models.timing import Timing
+from ingestify.domain.models.dataset.revision import RevisionState
+
+
+def JSONType(serializer=None, deserializer=None):
+    class _JsonType(TypeDecorator):
+        cache_ok = True
+        impl = JSON
+
+        def process_bind_param(self, value, dialect):
+            if serializer is not None:
+                return serializer(value)
+            return value
+
+        def process_result_value(self, value, dialect):
+            if deserializer is not None:
+                return deserializer(value)
+            return value
+
+    return _JsonType
 
 
 class TZDateTime(TypeDecorator):
@@ -25,7 +53,10 @@ class TZDateTime(TypeDecorator):
     LOCAL_TIMEZONE = datetime.datetime.utcnow().astimezone().tzinfo
     cache_ok = True
 
-    def process_bind_param(self, value: datetime, dialect):
+    def process_bind_param(self, value: Optional[datetime.datetime], dialect):
+        if not value:
+            return None
+
         if value.tzinfo is None:
             value = value.astimezone(self.LOCAL_TIMEZONE)
 
@@ -67,6 +98,45 @@ class DatasetStateString(TypeDecorator):
         return DatasetState[value]
 
 
+class RevisionStateString(TypeDecorator):
+    impl = String(255)
+
+    def process_bind_param(self, value: RevisionState, dialect):
+        return value.value
+
+    def process_result_value(self, value, dialect):
+        if not value:
+            return value
+
+        return RevisionState[value]
+
+
+class OperationString(TypeDecorator):
+    impl = String(255)
+
+    def process_bind_param(self, value: Operation, dialect):
+        return value.value
+
+    def process_result_value(self, value, dialect):
+        if not value:
+            return value
+
+        return Operation[value]
+
+
+class TaskStatusString(TypeDecorator):
+    impl = String(255)
+
+    def process_bind_param(self, value: TaskStatus, dialect):
+        return value.value
+
+    def process_result_value(self, value, dialect):
+        if not value:
+            return value
+
+        return TaskStatus[value]
+
+
 mapper_registry = registry()
 
 metadata = MetaData()
@@ -80,7 +150,7 @@ dataset_table = Table(
     Column("dataset_type", String(255)),
     Column("state", DatasetStateString),
     Column("name", String(255)),
-    Column("identifier", JSON),
+    Column("identifier", JSONType(deserializer=lambda item: Identifier(**item))),
     Column("metadata", JSON),
     Column("created_at", TZDateTime(6)),
     Column("updated_at", TZDateTime(6)),
@@ -95,7 +165,10 @@ revision_table = Table(
     Column("revision_id", Integer, primary_key=True),
     Column("description", String(255)),
     Column("created_at", TZDateTime(6)),
+    Column("state", RevisionStateString, default=RevisionState.PENDING_VALIDATION),
+    Column("source", JSONType()),
 )
+
 file_table = Table(
     "file",
     metadata,
@@ -129,7 +202,7 @@ mapper_registry.map_imperatively(
             Revision,
             backref="dataset",
             order_by=revision_table.c.revision_id,
-            lazy="joined",
+            lazy="selectin",
             cascade="all, delete-orphan",
         ),
     },
@@ -143,7 +216,7 @@ mapper_registry.map_imperatively(
             File,
             order_by=file_table.c.file_id,
             primaryjoin="and_(Revision.revision_id==File.revision_id, Revision.dataset_id==File.dataset_id)",
-            lazy="joined",
+            lazy="selectin",
             cascade="all, delete-orphan",
         )
     },
@@ -151,3 +224,110 @@ mapper_registry.map_imperatively(
 
 
 mapper_registry.map_imperatively(File, file_table)
+
+
+ingestion_job_summary = Table(
+    "ingestion_job_summary",
+    metadata,
+    Column("ingestion_job_id", String(255), primary_key=True),
+    # From the IngestionPlan
+    Column("source_name", String(255)),
+    Column("dataset_type", String(255)),
+    Column(
+        "data_spec_versions",
+        JSONType(
+            serializer=lambda data_spec_versions: data_spec_versions.to_dict(),
+            deserializer=lambda data_spec_versions: DataSpecVersionCollection.from_dict(
+                data_spec_versions
+            ),
+        ),
+    ),
+    Column(
+        "selector", JSONType(serializer=lambda selector: selector.filtered_attributes)
+    ),
+    Column("started_at", TZDateTime(6)),
+    Column("finished_at", TZDateTime(6)),
+    # Some task counters
+    Column("successful_tasks", Integer),
+    Column("ignored_successful_tasks", Integer),
+    Column("failed_tasks", Integer),
+    Column(
+        "timings",
+        JSONType(
+            serializer=lambda timings: [
+                timing.model_dump(mode="json") for timing in timings
+            ],
+            deserializer=lambda timings: [
+                Timing.model_validate(timing) for timing in timings
+            ],
+        ),
+    ),
+    # Column(
+    #     "task_summaries",
+    #     JSONType(
+    #         serializer=lambda task_summaries: [
+    #             task_summary.model_dump(mode="json") for task_summary in task_summaries
+    #         ],
+    #         deserializer=lambda task_summaries: [
+    #             TaskSummary.model_validate(task_summary)
+    #             for task_summary in task_summaries
+    #         ],
+    #     ),
+    # ),
+)
+
+
+task_summary_table = Table(
+    "task_summary",
+    metadata,
+    Column(
+        "ingestion_job_id",
+        String(255),
+        ForeignKey("ingestion_job_summary.ingestion_job_id"),
+        primary_key=True,
+    ),
+    Column("task_id", Integer, primary_key=True),
+    Column("started_at", TZDateTime(6)),
+    Column("ended_at", TZDateTime(6)),
+    Column("operation", OperationString),
+    Column(
+        "dataset_identifier", JSONType(deserializer=lambda item: Identifier(**item))
+    ),
+    Column("persisted_file_count", Integer),
+    Column("bytes_retrieved", Integer),
+    Column("last_modified", TZDateTime(6)),
+    Column("status", TaskStatusString),
+    Column(
+        "timings",
+        JSONType(
+            serializer=lambda timings: [
+                timing.model_dump(mode="json") for timing in timings
+            ],
+            deserializer=lambda timings: [
+                Timing.model_validate(timing) for timing in timings
+            ],
+        ),
+    ),
+    # Column("description", String(255)),
+    # Column("created_at", TZDateTime(6)),
+    # Column("state", RevisionStateString, default=RevisionState.PENDING_VALIDATION),
+    # Column("source", JSONType()),
+)
+
+
+mapper_registry.map_imperatively(
+    IngestionJobSummary,
+    ingestion_job_summary,
+    properties={
+        "task_summaries": relationship(
+            TaskSummary,
+            backref="ingestion_job_summary",
+            # order_by=task_summary_table.c.revision_id,
+            lazy="selectin",
+            cascade="all, delete-orphan",
+        ),
+    },
+)
+
+
+mapper_registry.map_imperatively(TaskSummary, task_summary_table)

ingestify/infra/store/dataset/sqlalchemy/repository.py

@@ -29,22 +29,6 @@ def parse_value(v):
         return v
 
 
-def json_serializer(o):
-    return json.dumps(o)
-
-
-def json_deserializer(o):
-    o = json.loads(o)
-    # THIS BREAKS WHEN USING OTHER JSON COLUMNS!!
-    o = Identifier(**o)
-    return o
-
-
-# @compiles(DateTime, "mysql")
-# def compile_datetime_mysql(type_, compiler, **kw):
-#     return "DATETIME(6)"
-
-
 def isfloat(x):
     try:
         a = float(x)
@@ -64,7 +48,7 @@ def isint(x):
         return a == b
 
 
-class SqlAlchemyDatasetRepository(DatasetRepository):
+class SqlAlchemySessionProvider:
     @staticmethod
     def fix_url(url: str) -> str:
         if url.startswith("postgres://"):
@@ -87,8 +71,6 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
             self.url,
             # Use the default isolation level, don't need SERIALIZABLE
             # isolation_level="SERIALIZABLE",
-            json_serializer=json_serializer,
-            json_deserializer=json_deserializer,
         )
         self.session = Session(bind=self.engine)
 
@@ -107,9 +89,29 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
         self.url = state["url"]
         self._init_engine()
 
+    def _close_engine(self):
+        if hasattr(self, "session"):
+            self.session.close()
+            self.engine.dispose()
+
     def __del__(self):
-        self.session.close()
-        self.engine.dispose()
+        self._close_engine()
+
+    def reset(self):
+        self._close_engine()
+        self._init_engine()
+
+    def get(self):
+        return self.session
+
+
+class SqlAlchemyDatasetRepository(DatasetRepository):
+    def __init__(self, session_provider: SqlAlchemySessionProvider):
+        self.session_provider = session_provider
+
+    @property
+    def session(self):
+        return self.session_provider.get()
 
     def _filter_query(
         self,
@@ -208,7 +210,7 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
 
         if not metadata_only:
             dataset_query = apply_query_filter(
-                self.session.query(Dataset).options(joinedload(Dataset.revisions))
+                self.session.query(Dataset)  # .options(joinedload(Dataset.revisions))
             )
             datasets = list(dataset_query)
         else:

ingestify/main.py

@@ -11,19 +11,20 @@ from ingestify import Source
 from ingestify.application.dataset_store import DatasetStore
 from ingestify.application.ingestion_engine import IngestionEngine
 from ingestify.application.secrets_manager import SecretsManager
-from ingestify.domain import Selector
-from ingestify.domain.models import (
-    dataset_repository_factory,
-    file_repository_factory,
-)
+from ingestify.domain import Selector, FileRepository
 from ingestify.domain.models.data_spec_version_collection import (
     DataSpecVersionCollection,
 )
 from ingestify.domain.models.event import EventBus, Publisher, Subscriber
 
-from ingestify.domain.models.extract_job import ExtractJob
+from ingestify.domain.models.ingestion.ingestion_plan import IngestionPlan
 from ingestify.domain.models.fetch_policy import FetchPolicy
 from ingestify.exceptions import ConfigurationError
+from ingestify.infra import S3FileRepository, LocalFileRepository
+from ingestify.infra.store.dataset.sqlalchemy import SqlAlchemyDatasetRepository
+from ingestify.infra.store.dataset.sqlalchemy.repository import (
+    SqlAlchemySessionProvider,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -59,8 +60,19 @@ def import_cls(name):
     return getattr(mod, components[-1])
 
 
+def build_file_repository(file_url: str) -> FileRepository:
+    if file_url.startswith("s3://"):
+        repository = S3FileRepository(url=file_url)
+    elif file_url.startswith("file://"):
+        repository = LocalFileRepository(url=file_url)
+    else:
+        raise Exception(f"Cannot find repository to handle file {file_url}")
+
+    return repository
+
+
 def get_dataset_store_by_urls(
-    dataset_url: str, file_url: str, bucket: str
+    metadata_url: str, file_url: str, bucket: str
 ) -> DatasetStore:
     """
     Initialize a DatasetStore by a DatasetRepository and a FileRepository
@@ -68,15 +80,18 @@ def get_dataset_store_by_urls(
     if not bucket:
         raise Exception("Bucket is not specified")
 
-    file_repository = file_repository_factory.build_if_supports(url=file_url)
+    file_repository = build_file_repository(file_url)
+
+    if secrets_manager.supports(metadata_url):
+        metadata_url = secrets_manager.load_as_db_url(metadata_url)
+
+    if metadata_url.startswith("postgres://"):
+        metadata_url = metadata_url.replace("postgress://", "postgress+")
 
-    if secrets_manager.supports(dataset_url):
-        dataset_url = secrets_manager.load_as_db_url(dataset_url)
+    sqlalchemy_session_provider = SqlAlchemySessionProvider(metadata_url)
 
-    if dataset_url.startswith("postgres://"):
-        dataset_url = dataset_url.replace("postgress://", "postgress+")
+    dataset_repository = SqlAlchemyDatasetRepository(sqlalchemy_session_provider)
 
-    dataset_repository = dataset_repository_factory.build_if_supports(url=dataset_url)
     return DatasetStore(
         dataset_repository=dataset_repository,
         file_repository=file_repository,
@@ -155,7 +170,7 @@ def get_engine(config_file, bucket: Optional[str] = None) -> IngestionEngine:
 
     logger.info("Initializing IngestionEngine")
     store = get_dataset_store_by_urls(
-        dataset_url=config["main"]["dataset_url"],
+        metadata_url=config["main"]["metadata_url"],
         file_url=config["main"]["file_url"],
         bucket=bucket or config["main"].get("default_bucket"),
     )
@@ -177,15 +192,20 @@ def get_engine(config_file, bucket: Optional[str] = None) -> IngestionEngine:
 
     fetch_policy = FetchPolicy()
 
-    for job in config["extract_jobs"]:
+    # Previous naming
+    ingestion_plans = config.get("extract_jobs", [])
+    # New naming
+    ingestion_plans.extend(config.get("ingestion_plans", []))
+
+    for ingestion_plan in ingestion_plans:
         data_spec_versions = DataSpecVersionCollection.from_dict(
-            job.get("data_spec_versions", {"default": {"v1"}})
+            ingestion_plan.get("data_spec_versions", {"default": {"v1"}})
         )
 
-        if "selectors" in job:
+        if "selectors" in ingestion_plan:
             selectors = [
                 Selector.build(selector, data_spec_versions=data_spec_versions)
-                for selector_args in job["selectors"]
+                for selector_args in ingestion_plan["selectors"]
                 for selector in _product_selectors(selector_args)
             ]
         else:
@@ -193,13 +213,13 @@ def get_engine(config_file, bucket: Optional[str] = None) -> IngestionEngine:
             # but makes it easier later one where we loop over selectors.
             selectors = [Selector.build({}, data_spec_versions=data_spec_versions)]
 
-        import_job = ExtractJob(
-            source=sources[job["source"]],
-            dataset_type=job["dataset_type"],
+        ingestion_plan = IngestionPlan(
+            source=sources[ingestion_plan["source"]],
+            dataset_type=ingestion_plan["dataset_type"],
             selectors=selectors,
             fetch_policy=fetch_policy,
             data_spec_versions=data_spec_versions,
         )
-        ingestion_engine.add_extract_job(import_job)
+        ingestion_engine.add_ingestion_plan(ingestion_plan)
 
     return ingestion_engine