ingestify 0.4.0__tar.gz → 0.4.2__tar.gz

{ingestify-0.4.0 → ingestify-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ingestify
-Version: 0.4.0
+Version: 0.4.2
 Summary: Data Ingestion Framework
 Author: Koen Vossen
 Author-email: info@koenvossen.nl

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/__init__.py

@@ -8,4 +8,4 @@ if not __INGESTIFY_SETUP__:
     from .infra import retrieve_http
     from .source_base import Source, DatasetResource
 
-__version__ = "0.4.0"
+__version__ = "0.4.2"

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/application/loader.py

@@ -35,11 +35,8 @@ class Loader:
         provider: Optional[str] = None,
         source: Optional[str] = None,
     ):
-        # First collect all selectors, before discovering datasets
-        selectors = {}
+        ingestion_plans = []
         for ingestion_plan in self.ingestion_plans:
-            logger.info(f"Determining selectors for {ingestion_plan}")
-
             if provider is not None:
                 if ingestion_plan.source.provider != provider:
                     logger.info(
@@ -54,6 +51,13 @@ class Loader:
                     )
                     continue
 
+            ingestion_plans.append(ingestion_plan)
+
+        # First collect all selectors, before discovering datasets
+        selectors = {}
+        for ingestion_plan in ingestion_plans:
+            logger.info(f"Determining selectors for {ingestion_plan}")
+
             static_selectors = [
                 selector
                 for selector in ingestion_plan.selectors

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/cmdline.py

@@ -88,6 +88,14 @@ def cli():
     help="Source - only run tasks for a single source",
     type=str,
 )
+@click.option(
+    "--disable-events",
+    "disable_events",
+    required=False,
+    help="Disable events - disable all event handlers",
+    is_flag=True,
+    type=bool,
+)
 def run(
     config_file: str,
     bucket: Optional[str],
@@ -95,9 +103,10 @@ def run(
     provider: Optional[str],
     source: Optional[str],
     debug: Optional[bool],
+    disable_events: Optional[bool],
 ):
     try:
-        engine = get_engine(config_file, bucket)
+        engine = get_engine(config_file, bucket, disable_events=disable_events)
     except ConfigurationError as e:
         if debug:
             raise

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/domain/models/ingestion/ingestion_job.py

@@ -218,7 +218,7 @@ class IngestionJob:
         # Process all items in batches. Yield a IngestionJobSummary per batch
 
         logger.info("Finding metadata")
-        with ingestion_job_summary.record_timing("get_dataset_collection"):
+        with ingestion_job_summary.record_timing("get_dataset_collection_metadata"):
             dataset_collection_metadata = store.get_dataset_collection(
                 dataset_type=self.ingestion_plan.dataset_type,
                 provider=self.ingestion_plan.source.provider,
@@ -232,6 +232,7 @@ class IngestionJob:
         # 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
         try:
+            logger.info(f"Finding datasets for selector={self.selector}")
             with ingestion_job_summary.record_timing("find_datasets"):
                 dataset_resources = self.ingestion_plan.source.find_datasets(
                     dataset_type=self.ingestion_plan.dataset_type,
@@ -249,6 +250,8 @@ class IngestionJob:
             yield ingestion_job_summary
             return
 
+        logger.info("Starting tasks")
+
         finish_task_timer = ingestion_job_summary.start_timing("tasks")
 
         while True:
@@ -273,13 +276,16 @@ class IngestionJob:
                 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,
-            )
+            logger.info(f"Searching for existing Datasets for DatasetResources")
+
+            with ingestion_job_summary.record_timing("get_dataset_collection"):
+                # 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
 

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/domain/models/ingestion/ingestion_job_summary.py

@@ -66,7 +66,7 @@ class IngestionJobSummary(BaseModel, HasTiming):
         self.skipped_datasets += skipped_datasets
 
     def task_count(self):
-        return len(self.task_summaries)
+        return len(self.task_summaries) + self.skipped_datasets
 
     def _set_ended(self):
         self.failed_tasks = len(
@@ -106,22 +106,22 @@ class IngestionJobSummary(BaseModel, HasTiming):
             f"\nIngestionJobSummary {self.state.value} 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: ")
+        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"   - {timing.name}: {format_duration(timing.duration)}")
         print(
-            f"*  - Tasks: {len(self.task_summaries)} - {(len(self.task_summaries) / self.duration.total_seconds()):.1f} tasks/sec"
+            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.ignored_successful_tasks}")
-        print(f"*    - Skipped datasets: {self.skipped_datasets}")
+        print(f"   - Failed tasks: {self.failed_tasks}")
+        print(f"   - Successful tasks: {self.successful_tasks}")
+        print(f"   - Successful ignored tasks: {self.ignored_successful_tasks}")
+        print(f"   - Skipped datasets: {self.skipped_datasets}")
         print("********************************")
 
     def __enter__(self):

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/infra/store/dataset/sqlalchemy/repository.py

@@ -1,32 +1,33 @@
 import itertools
-import json
+import logging
 import uuid
-from collections import defaultdict
 from typing import Optional, Union, List
 
 from sqlalchemy import (
     create_engine,
     func,
     text,
-    tuple_,
     Table,
-    insert,
-    Transaction,
     Connection,
+    union_all,
+    literal,
+    select,
+    and_,
+    Column,
+    or_,
+    Dialect,
 )
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import NoSuchModuleError
-from sqlalchemy.orm import Session, joinedload
+from sqlalchemy.orm import Session, Query
 
 from ingestify.domain import File, Revision
 from ingestify.domain.models import (
     Dataset,
     DatasetCollection,
     DatasetRepository,
-    Identifier,
     Selector,
 )
-from ingestify.domain.models.base import BaseModel
 from ingestify.domain.models.dataset.collection_metadata import (
     DatasetCollectionMetadata,
 )
@@ -43,6 +44,8 @@ from .tables import (
     task_summary_table,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def parse_value(v):
     try:
@@ -94,6 +97,7 @@ class SqlAlchemySessionProvider:
             # Use the default isolation level, don't need SERIALIZABLE
             # isolation_level="SERIALIZABLE",
         )
+        self.dialect = self.engine.dialect
         self.session = Session(bind=self.engine)
 
     def __init__(self, url: str):
@@ -111,22 +115,26 @@ class SqlAlchemySessionProvider:
         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._close_engine()
+        self.close()
 
     def reset(self):
-        self._close_engine()
+        self.close()
         self._init_engine()
 
+    def close(self):
+        if hasattr(self, "session"):
+            self.session.close()
+            self.engine.dispose()
+
     def get(self):
         return self.session
 
 
+def in_(column: Column, values):
+    return or_(*[column == value for value in values])
+
+
 class SqlAlchemyDatasetRepository(DatasetRepository):
     def __init__(self, session_provider: SqlAlchemySessionProvider):
         self.session_provider = session_provider
@@ -135,8 +143,12 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
     def session(self):
         return self.session_provider.get()
 
+    @property
+    def dialect(self) -> Dialect:
+        return self.session_provider.dialect
+
     def _upsert(self, connection: Connection, table: Table, entities: list[dict]):
-        dialect = self.session.bind.dialect.name
+        dialect = self.dialect.name
         if dialect == "mysql":
             from sqlalchemy.dialects.mysql import insert
         elif dialect == "postgresql":
@@ -169,11 +181,6 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
         dataset_id: Optional[Union[str, List[str]]] = None,
         selector: Optional[Union[Selector, List[Selector]]] = None,
     ):
-        query = query.filter(dataset_table.c.bucket == bucket)
-        if dataset_type:
-            query = query.filter(dataset_table.c.dataset_type == dataset_type)
-        if provider:
-            query = query.filter(dataset_table.c.provider == provider)
         if dataset_id is not None:
             if isinstance(dataset_id, list):
                 if len(dataset_id) == 0:
@@ -181,11 +188,11 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
                     # return an empty DatasetCollection
                     return DatasetCollection()
 
-                query = query.filter(dataset_table.c.dataset_id.in_(dataset_id))
+                query = query.filter(in_(dataset_table.c.dataset_id, dataset_id))
             else:
                 query = query.filter(dataset_table.c.dataset_id == dataset_id)
 
-        dialect = self.session.bind.dialect.name
+        dialect = self.dialect.name
 
         if not isinstance(selector, list):
             where, selector = selector.split("where")
@@ -201,13 +208,22 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
             if not selectors:
                 raise ValueError("Selectors must contain at least one item")
 
-            keys = list(selectors[0].filtered_attributes.keys())
+            attribute_sets = {
+                tuple(selector.filtered_attributes.items()) for selector in selectors
+            }
+
+            # Define a virtual table using a CTE for all attributes
+            attribute_cte = union_all(
+                *[
+                    select(*(literal(value).label(key) for key, value in attr_set))
+                    for attr_set in attribute_sets
+                ]
+            ).cte("attributes")
 
-            columns = []
+            keys = list(selectors[0].filtered_attributes.keys())
             first_selector = selectors[0].filtered_attributes
 
-            # Create a query like this:
-            #  SELECT * FROM dataset WHERE (column1, column2, column3) IN ((1, 2, 3), (4, 5, 6), (7, 8, 9))
+            join_conditions = []
             for k in keys:
                 if dialect == "postgresql":
                     column = dataset_table.c.identifier[k]
@@ -215,40 +231,43 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
                     # Take the value from the first selector to determine the type.
                     # TODO: check all selectors to determine the type
                     v = first_selector[k]
-                    if isint(v):
+                    if isinstance(v, int):
                         column = column.as_integer()
-                    elif isfloat(v):
-                        column = column.as_float()
                     else:
                         column = column.as_string()
                 else:
                     column = func.json_extract(dataset_table.c.identifier, f"$.{k}")
-                columns.append(column)
 
-            values = []
-            for selector in selectors:
-                filtered_attributes = selector.filtered_attributes
-                values.append(tuple([filtered_attributes[k] for k in keys]))
+                join_conditions.append(attribute_cte.c[k] == column)
 
-            query = query.filter(tuple_(*columns).in_(values))
+            query = query.select_from(
+                dataset_table.join(attribute_cte, and_(*join_conditions))
+            )
 
         if where:
             query = query.filter(text(where))
+
+        query = query.filter(dataset_table.c.bucket == bucket)
+        if dataset_type:
+            query = query.filter(dataset_table.c.dataset_type == dataset_type)
+        if provider:
+            query = query.filter(dataset_table.c.provider == provider)
+
         return query
 
-    def load_datasets(self, dataset_ids: list[str]) -> list[Dataset]:
+    def _load_datasets(self, dataset_ids: list[str]) -> list[Dataset]:
         if not dataset_ids:
             return []
 
         dataset_rows = list(
             self.session.query(dataset_table).filter(
-                dataset_table.c.dataset_id.in_(dataset_ids)
+                in_(dataset_table.c.dataset_id, dataset_ids)
             )
         )
         revisions_per_dataset = {}
         rows = (
             self.session.query(revision_table)
-            .filter(revision_table.c.dataset_id.in_(dataset_ids))
+            .filter(in_(revision_table.c.dataset_id, dataset_ids))
             .order_by(revision_table.c.dataset_id)
         )
 
@@ -260,7 +279,7 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
         files_per_revision = {}
         rows = (
             self.session.query(file_table)
-            .filter(file_table.c.dataset_id.in_(dataset_ids))
+            .filter(in_(file_table.c.dataset_id, dataset_ids))
             .order_by(file_table.c.dataset_id, file_table.c.revision_id)
         )
 
@@ -290,6 +309,12 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
             )
         return datasets
 
+    def _debug_query(self, q: Query):
+        text_ = q.statement.compile(
+            compile_kwargs={"literal_binds": True}, dialect=self.dialect
+        )
+        logger.debug(f"Running query: {text_}")
+
     def get_dataset_collection(
         self,
         bucket: str,
@@ -309,22 +334,40 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
                 selector=selector,
             )
 
-        if not metadata_only:
-            dataset_query = apply_query_filter(
-                self.session.query(dataset_table.c.dataset_id)
-            )
-            dataset_ids = [row.dataset_id for row in dataset_query]
-            datasets = self.load_datasets(dataset_ids)
-        else:
-            datasets = []
+        with self.session:
+            # Use a contextmanager to make sure it's closed afterwards
 
-        metadata_result_row = apply_query_filter(
-            self.session.query(
-                func.max(dataset_table.c.last_modified_at).label("last_modified_at"),
-                func.count().label("row_count"),
-            )
-        ).first()
-        dataset_collection_metadata = DatasetCollectionMetadata(*metadata_result_row)
+            if not metadata_only:
+                dataset_query = apply_query_filter(
+                    self.session.query(dataset_table.c.dataset_id)
+                )
+                self._debug_query(dataset_query)
+                dataset_ids = [row.dataset_id for row in dataset_query]
+                datasets = self._load_datasets(dataset_ids)
+
+                dataset_collection_metadata = DatasetCollectionMetadata(
+                    last_modified=max(dataset.last_modified_at for dataset in datasets)
+                    if datasets
+                    else None,
+                    row_count=len(datasets),
+                )
+            else:
+                datasets = []
+
+                metadata_result_query = apply_query_filter(
+                    self.session.query(
+                        func.max(dataset_table.c.last_modified_at).label(
+                            "last_modified_at"
+                        ),
+                        func.count().label("row_count"),
+                    )
+                )
+
+                self._debug_query(metadata_result_query)
+
+                dataset_collection_metadata = DatasetCollectionMetadata(
+                    *metadata_result_query.first()
+                )
 
         return DatasetCollection(dataset_collection_metadata, datasets)
 
@@ -337,6 +380,9 @@ class SqlAlchemyDatasetRepository(DatasetRepository):
     def connect(self):
         return self.session_provider.engine.connect()
 
+    def __del__(self):
+        self.session_provider.close()
+
     def _save(self, datasets: list[Dataset]):
         """Only do upserts. Never delete. Rows get only deleted when an entire Dataset is removed."""
         datasets_entities = []

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/infra/store/dataset/sqlalchemy/tables.py

@@ -14,6 +14,7 @@ from sqlalchemy import (
     String,
     Table,
     TypeDecorator,
+    Index,
 )
 
 from sqlalchemy.dialects.postgresql import JSONB
@@ -167,6 +168,15 @@ dataset_table = Table(
     Column("created_at", TZDateTime(6)),
     Column("updated_at", TZDateTime(6)),
     Column("last_modified_at", TZDateTime(6)),
+    # Required for performance querying when there are a lot of Datasets
+    # with the same provider and dataset_type
+    Index(
+        "idx_bucket_type_provider_last_modified",
+        "bucket",
+        "provider",
+        "dataset_type",
+        "last_modified_at",
+    ),
 )
 
 revision_table = Table(

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify/main.py

@@ -182,7 +182,9 @@ def get_event_subscriber_cls(key: str) -> Type[Subscriber]:
     return import_cls(key)
 
 
-def get_engine(config_file, bucket: Optional[str] = None) -> IngestionEngine:
+def get_engine(
+    config_file, bucket: Optional[str] = None, disable_events: bool = False
+) -> IngestionEngine:
     config = parse_config(config_file, default_value="")
 
     logger.info("Initializing sources")
@@ -201,11 +203,16 @@ def get_engine(config_file, bucket: Optional[str] = None) -> IngestionEngine:
 
     # Setup an EventBus and wire some more components
     event_bus = EventBus()
-    publisher = Publisher()
-    for subscriber in config.get("event_subscribers", []):
-        cls = get_event_subscriber_cls(subscriber["type"])
-        publisher.add_subscriber(cls(store))
-    event_bus.register(publisher)
+    if not disable_events:
+        # When we disable all events we don't register any publishers
+        publisher = Publisher()
+        for subscriber in config.get("event_subscribers", []):
+            cls = get_event_subscriber_cls(subscriber["type"])
+            publisher.add_subscriber(cls(store))
+        event_bus.register(publisher)
+    else:
+        logger.info("Disabling all event handlers")
+
     store.set_event_bus(event_bus)
 
     ingestion_engine = IngestionEngine(

{ingestify-0.4.0 → ingestify-0.4.2}/ingestify.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ingestify
-Version: 0.4.0
+Version: 0.4.2
 Summary: Data Ingestion Framework
 Author: Koen Vossen
 Author-email: info@koenvossen.nl