neo4j-etl-lib 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

etl_lib/__init__.py

@@ -1,4 +1,4 @@
 """
 Building blocks for ETL pipelines.
 """
-__version__ = "0.1.1"
+__version__ = "0.3.0"

etl_lib/cli/run_tools.py

@@ -67,7 +67,7 @@ def __driver(ctx):
 @click.pass_context
 def cli(ctx, neo4j_uri, neo4j_user, neo4j_password, log_file, database_name):
     """
-        Command-line tool to process files in INPUT_DIRECTORY.
+        Command-line tool for ETL pipelines.
 
         Environment variables can be configured via a .env file or overridden via CLI options:
 

etl_lib/core/BatchProcessor.py

@@ -2,7 +2,7 @@ import abc
 import logging
 import sys
 from dataclasses import dataclass, field
-from typing import Generator
+from typing import Generator, List, Any
 
 from etl_lib.core.Task import Task
 from etl_lib.core.utils import merge_summery
@@ -13,7 +13,7 @@ class BatchResults:
     """
     Return object of the :py:func:`~BatchProcessor.get_batch` method, wrapping a batched data together with meta information.
     """
-    chunk: []
+    chunk: List[Any]
     """The batch of data."""
     statistics: dict = field(default_factory=dict)
     """`dict` of statistic information, such as row processed, nodes writen, .."""
@@ -38,11 +38,11 @@ def append_result(org: BatchResults, stats: dict) -> BatchResults:
                         batch_size=org.batch_size)
 
 
-class BatchProcessor:
+class BatchProcessor(abc.ABC):
     """
     Allows assembly of :py:class:`etl_lib.core.Task.Task` out of smaller building blocks.
 
-    This way, functionally such as reading from a CSV file, writing to a database or validation
+    This way, functionally, such as reading from a CSV file, writing to a database or validation
     can be implemented and tested independently and re-used.
 
     BatchProcessors form, a linked list, where each processor only knows about its predecessor.
@@ -57,17 +57,17 @@ class BatchProcessor:
         Constructs a new :py:class:`etl_lib.core.BatchProcessor` instance.
 
         Args:
-            context: :py:class:`etl_lib.core.ETLContext.ETLContext` instance. Will be available to subclasses.
+            context: :py:class:`etl_lib.core.ETLContext.ETLContext` instance. It Will be available to subclasses.
             task: :py:class:`etl_lib.core.Task.Task` this processor is part of.
                 Needed for status reporting only.
             predecessor: Source of batches for this processor.
-                Can be `None` of no predecessor is needed (such as when this processor is the start of the queue.
+                Can be `None` if no predecessor is needed (such as when this processor is the start of the queue).
         """
         self.context = context
         """:py:class:`etl_lib.core.ETLContext.ETLContext` instance. Providing access to general facilities."""
         self.predecessor = predecessor
         """Predecessor, used as a source of batches."""
-        self.logger = logging.getLogger(self.__class__.__name__)
+        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
         self.task = task
         """The :py:class:`etl_lib.core.Task.Task` owning instance."""
 

etl_lib/core/ClosedLoopBatchProcessor.py

@@ -1,7 +1,7 @@
 from typing import Generator
 
-from etl_lib.core.ETLContext import ETLContext
 from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults, append_result
+from etl_lib.core.ETLContext import ETLContext
 from etl_lib.core.Task import Task
 
 
@@ -24,7 +24,13 @@ class ClosedLoopBatchProcessor(BatchProcessor):
         for batch in self.predecessor.get_batch(max_batch__size):
             result = append_result(result, batch.statistics)
             batch_cnt += 1
-            self.context.reporter.report_progress(self.task, batch_cnt, self.expected_rows, result.statistics)
+            self.context.reporter.report_progress(self.task, batch_cnt, self._safe_calculate_count(max_batch__size),
+                                                  result.statistics)
 
         self.logger.debug(result.statistics)
         yield result
+
+    def _safe_calculate_count(self, batch_size: int | None) -> int:
+        if not self.expected_rows or not batch_size:
+            return 0
+        return (self.expected_rows + batch_size - 1) // batch_size  # ceiling division

etl_lib/core/ETLContext.py

@@ -1,22 +1,58 @@
 import logging
-from typing import NamedTuple, Any
+from typing import Any, Dict, List, NamedTuple
 
-from graphdatascience import GraphDataScience
-from neo4j import GraphDatabase, WRITE_ACCESS, SummaryCounters
+from neo4j.exceptions import Neo4jError
+
+try:
+    from graphdatascience import GraphDataScience
+
+    gds_available = False
+except ImportError:
+    gds_available = False
+    logging.info("Graph Data Science not installed, skipping")
+    GraphDataScience = None
+
+from neo4j import GraphDatabase, Session, WRITE_ACCESS, SummaryCounters
+
+try:
+    from sqlalchemy import create_engine
+    from sqlalchemy.engine import Engine
+
+    sqlalchemy_available = True
+except ImportError:
+    sqlalchemy_available = False
+    logging.info("SQL Alchemy not installed, skipping")
+    create_engine = None  # this and next line needed to prevent PyCharm warning
+    Engine = None
 
 from etl_lib.core.ProgressReporter import get_reporter
 
 
 class QueryResult(NamedTuple):
     """Result of a query against the neo4j database."""
-    data: []
+    data: List[Any]
     """Data as returned from the query."""
-    summery: {}
+    summery: Dict[str, int]
     """Counters as reported by neo4j. Contains entries such as `nodes_created`, `nodes_deleted`, etc."""
 
 
 def append_results(r1: QueryResult, r2: QueryResult) -> QueryResult:
-    return QueryResult(r1.data + r2.data, r1.summery + r2.summery)
+    """
+    Appends two QueryResult objects, summing the values for duplicate keys in the summary.
+
+    Args:
+        r1: The first QueryResult object.
+        r2: The second QueryResult object to append.
+
+    Returns:
+        A new QueryResult object with combined data and summed summary counts.
+    """
+    combined_summery = r1.summery.copy()
+
+    for key, value in r2.summery.items():
+        combined_summery[key] = combined_summery.get(key, 0) + value
+
+    return QueryResult(r1.data + r2.data, combined_summery)
 
 
 class Neo4jContext:
@@ -34,37 +70,38 @@ class Neo4jContext:
         - `NEO4J_PASSWORD`.
         - `NEO4J_DATABASE`,
         """
-        self.logger = logging.getLogger(self.__class__.__name__)
+        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
         self.uri = env_vars["NEO4J_URI"]
         self.auth = (env_vars["NEO4J_USERNAME"],
                      env_vars["NEO4J_PASSWORD"])
         self.database = env_vars["NEO4J_DATABASE"]
         self.__neo4j_connect()
 
-    def query_database(self, session, query, **kwargs) -> QueryResult:
+    def query_database(self, session: Session, query, **kwargs) -> QueryResult:
         """
-        Executes a Cypher query on the Neo4j database.
-
-        Args:
-            session: Neo4j database session.
-            query: Cypher query either as a single query or as a list.
+        Executes Cypher and returns (records, counters) with retryable write semantics.
+        Accepts either a single query string or a list of queries.
+        Does not work with CALL {} IN TRANSACTION queries.
         """
         if isinstance(query, list):
-            results = []
-            for single_query in query:
-                result = self.query_database(session, single_query, **kwargs)
-                results = append_results(results, result)
+            results = None
+            for single in query:
+                part = self.query_database(session, single, **kwargs)
+                results = append_results(results, part) if results is not None else part
             return results
-        else:
-            try:
-                res = session.run(query, **kwargs)
-                counters = res.consume().counters
 
-                return QueryResult(res, self.__counters_2_dict(counters))
+        def _tx(tx, q, params):
+            res = tx.run(q, **params)
+            records = list(res)
+            counters = res.consume().counters
+            return records, counters
 
-            except Exception as e:
-                self.logger.error(e)
-                raise e
+        try:
+            records, counters = session.execute_write(_tx, query, kwargs)
+            return QueryResult(records, self.__counters_2_dict(counters))
+        except Neo4jError as e:
+            self.logger.error(e)
+            raise
 
     @staticmethod
     def __counters_2_dict(counters: SummaryCounters):
@@ -99,22 +136,6 @@ class Neo4jContext:
         else:
             return self.driver.session(database=database, default_access_mode=WRITE_ACCESS)
 
-    def gds(self, database=None) -> GraphDataScience:
-        """
-        Creates a new GraphDataScience client.
-
-        Args:
-            database: Name of the database to use for this dgs client.
-                If not provided, the database name provided during construction will be used.
-
-        Returns:
-            gds client.
-        """
-        if database is None:
-            return GraphDataScience.from_neo4j_driver(driver=self.driver, database=self.database)
-        else:
-            return GraphDataScience.from_neo4j_driver(driver=self.driver, database=database)
-
     def __neo4j_connect(self):
         self.driver = GraphDatabase.driver(uri=self.uri, auth=self.auth,
                                            notifications_min_severity="OFF")
@@ -123,11 +144,50 @@ class Neo4jContext:
             f"driver connected to instance at {self.uri} with username {self.auth[0]} and database {self.database}")
 
 
+def gds(neo4j_context) -> GraphDataScience:
+    """
+    Creates a new GraphDataScience client.
+
+    Args:
+        neo4j_context: Neo4j context containing driver and database name.
+
+    Returns:
+        gds client.
+    """
+    return GraphDataScience.from_neo4j_driver(driver=neo4j_context.driver, database=neo4j_context.database)
+
+
+if sqlalchemy_available:
+    class SQLContext:
+        def __init__(self, database_url: str, pool_size: int = 10, max_overflow: int = 20):
+            """
+            Initializes the SQL context with an SQLAlchemy engine.
+
+            Args:
+                database_url (str): SQLAlchemy connection URL.
+                pool_size (int): Number of connections to maintain in the pool.
+                max_overflow (int): Additional connections allowed beyond pool_size.
+            """
+            self.engine: Engine = create_engine(
+                database_url,
+                pool_pre_ping=True,
+                pool_size=pool_size,
+                max_overflow=max_overflow,
+                pool_recycle=1800,  # recycle connections older than 30m
+                connect_args={
+                    # turn on TCP keepalives on the client socket:
+                    "keepalives": 1,
+                    "keepalives_idle": 60,  # after 60s of idle
+                    "keepalives_interval": 10,  # probe every 10s
+                    "keepalives_count": 5,  # give up after 5 failed probes
+                })
+
+
 class ETLContext:
     """
     General context information.
 
-    Will be passed to all :py:class:`etl_lib.core.Task` to provide access to environment variables and functionally
+    Will be passed to all :class:`~etl_lib.core.Task.Task` to provide access to environment variables and functionally
     deemed general enough that all parts of the ETL pipeline would need it.
     """
 
@@ -136,15 +196,20 @@ class ETLContext:
         Create a new ETLContext.
 
         Args:
-            env_vars: Environment variables. Stored internally and can be accessed via :py:func:`~env` .
+            env_vars: Environment variables. Stored internally and can be accessed via :func:`~env` .
 
-        The context created will contain an :py:class:`~Neo4jContext` and a :py:class:`ProgressReporter`.
+        The context created will contain an :class:`~Neo4jContext` and a :class:`~etl_lib.core.ProgressReporter.ProgressReporter`.
         See there for keys used from the provided `env_vars` dict.
         """
-        self.logger = logging.getLogger(self.__class__.__name__)
+        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
         self.neo4j = Neo4jContext(env_vars)
         self.__env_vars = env_vars
         self.reporter = get_reporter(self)
+        sql_uri = self.env("SQLALCHEMY_URI")
+        if sql_uri is not None and sqlalchemy_available:
+            self.sql = SQLContext(sql_uri)
+        if gds_available:
+            self.gds = gds(self.neo4j)
 
     def env(self, key: str) -> Any:
         """
@@ -154,7 +219,8 @@ class ETLContext:
             key: name of the entry to read.
 
         Returns:
-            va  lue of the entry, or None if the key is not in the dict.
+            value of the entry, or None if the key is not in the dict.
         """
         if key in self.__env_vars:
             return self.__env_vars[key]
+        return None

etl_lib/core/ParallelBatchProcessor.py

@@ -0,0 +1,180 @@
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Generator, List
+
+from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults
+from etl_lib.core.utils import merge_summery
+
+
+class ParallelBatchResult(BatchResults):
+    """
+    Represents a batch split into parallelizable partitions.
+
+    `chunk` is a list of lists, each sub-list is a partition.
+    """
+    pass
+
+
+class ParallelBatchProcessor(BatchProcessor):
+    """
+    BatchProcessor that runs worker threads over partitions of batches.
+
+    Receives a special BatchResult (:py:class:`ParallelBatchResult`) from the predecessor.
+    All chunks in a ParallelBatchResult it receives can be processed in parallel.
+    See :py:class:`etl_lib.core.SplittingBatchProcessor` on how to produce them.
+    Prefetches the next ParallelBatchResults from its predecessor.
+    The actual processing of the batches is deferred to the configured worker.
+
+    Note:
+        - The predecessor must emit `ParallelBatchResult` instances.
+
+    Args:
+        context: ETL context.
+        worker_factory: A zero-arg callable that returns a new BatchProcessor
+                        each time it's called. This processor is responsible for the processing pf the batches.
+        task: optional Task for reporting.
+        predecessor: upstream BatchProcessor that must emit ParallelBatchResult.
+        max_workers: number of parallel threads for partitions.
+        prefetch: number of ParallelBatchResults to prefetch from the predecessor.
+
+    Behavior:
+        - For every ParallelBatchResult, spins up `max_workers` threads.
+        - Each thread calls its own worker from `worker_factory()`, with its
+          partition wrapped by `SingleBatchWrapper`.
+        - Collects and merges their BatchResults in a fail-fast manner: on first
+          exception, logs the error, cancels remaining threads, and raises an exception.
+    """
+
+    def __init__(
+            self,
+            context,
+            worker_factory: Callable[[], BatchProcessor],
+            task=None,
+            predecessor=None,
+            max_workers: int = 4,
+            prefetch: int = 4,
+    ):
+        super().__init__(context, task, predecessor)
+        self.worker_factory = worker_factory
+        self.max_workers = max_workers
+        self.prefetch = prefetch
+        self._batches_done = 0
+
+    def _process_parallel(self, pbr: ParallelBatchResult) -> BatchResults:
+        """
+        Run one worker per partition in `pbr.chunk`, merge their outputs, and include upstream
+        statistics from `pbr.statistics` so counters (e.g., valid/invalid rows from validation)
+        are preserved through the parallel stage.
+
+        Progress reporting:
+          - After each partition completes, report batch count only
+        """
+        merged_stats = dict(pbr.statistics or {})
+        merged_chunk = []
+        total = 0
+
+        parts_total = len(pbr.chunk)
+        partitions_done = 0
+
+        self.logger.debug(f"Processing pbr of len {parts_total}")
+        with ThreadPoolExecutor(max_workers=self.max_workers, thread_name_prefix='PBP_worker_') as pool:
+            futures = [pool.submit(self._process_partition, part) for part in pbr.chunk]
+            try:
+                for f in as_completed(futures):
+                    out = f.result()
+
+                    # Merge into this PBR's cumulative result (returned downstream)
+                    merged_stats = merge_summery(merged_stats, out.statistics or {})
+                    total += out.batch_size
+                    merged_chunk.extend(out.chunk if isinstance(out.chunk, list) else [out.chunk])
+
+                    partitions_done += 1
+                    self.context.reporter.report_progress(
+                        task=self.task,
+                        batches=self._batches_done,
+                        expected_batches=None,
+                        stats={},
+                    )
+
+            except Exception as e:
+                for g in futures:
+                    g.cancel()
+                pool.shutdown(cancel_futures=True)
+                raise RuntimeError("partition processing failed") from e
+
+        self.logger.debug(f"Finished processing pbr with {merged_stats}")
+        return BatchResults(chunk=merged_chunk, statistics=merged_stats, batch_size=total)
+
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
+        """
+        Pulls ParallelBatchResult batches from the predecessor, prefetching
+        up to `prefetch` ahead, processes each batch's partitions in
+        parallel threads, and yields a flattened BatchResults. The predecessor
+        can run ahead while the current batch is processed.
+        """
+        pbr_queue: queue.Queue[ParallelBatchResult | object] = queue.Queue(self.prefetch)
+        SENTINEL = object()
+        exc: BaseException | None = None
+
+        def producer():
+            nonlocal exc
+            try:
+                for pbr in self.predecessor.get_batch(max_batch_size):
+                    self.logger.debug(
+                        f"adding pgr {pbr.statistics} / {len(pbr.chunk)} to queue of size {pbr_queue.qsize()}"
+                    )
+                    pbr_queue.put(pbr)
+            except BaseException as e:
+                exc = e
+            finally:
+                pbr_queue.put(SENTINEL)
+
+        threading.Thread(target=producer, daemon=True, name='prefetcher').start()
+
+        while True:
+            pbr = pbr_queue.get()
+            if pbr is SENTINEL:
+                if exc is not None:
+                    self.logger.error("Upstream producer failed", exc_info=True)
+                    raise exc
+                break
+            result = self._process_parallel(pbr)
+            yield result
+
+    class SingleBatchWrapper(BatchProcessor):
+        """
+        Simple BatchProcessor that returns the batch it receives via init.
+        Will be used as predecessor for the worker
+        """
+
+        def __init__(self, context, batch: List[Any]):
+            super().__init__(context=context, predecessor=None)
+            self._batch = batch
+
+        def get_batch(self, max_batch__size: int) -> Generator[BatchResults, None, None]:
+            # Ignores max_size; yields exactly one BatchResults containing the whole batch
+            yield BatchResults(
+                chunk=self._batch,
+                statistics={},
+                batch_size=len(self._batch)
+            )
+
+    def _process_partition(self, partition):
+        """
+        Processes one partition of items by:
+          1. Wrapping it in SingleBatchWrapper
+          2. Instantiating a fresh worker via worker_factory()
+          3. Setting the worker's predecessor to the wrapper
+          4. Running exactly one batch and returning its BatchResults
+
+        Raises whatever exception the worker raises, allowing _process_parallel
+        to handle fail-fast behavior.
+        """
+        self.logger.debug("Processing partition")
+        wrapper = self.SingleBatchWrapper(self.context, partition)
+        worker = self.worker_factory()
+        worker.predecessor = wrapper
+        result = next(worker.get_batch(len(partition)))
+        self.logger.debug(f"finished processing partition with {result.statistics}")
+        return result

etl_lib/core/ProgressReporter.py

@@ -4,6 +4,7 @@ from datetime import datetime
 from tabulate import tabulate
 
 from etl_lib.core.Task import Task, TaskGroup, TaskReturn
+from etl_lib.core.utils import add_sigint_handler
 
 
 class ProgressReporter:
@@ -18,7 +19,7 @@ class ProgressReporter:
 
     def __init__(self, context):
         self.context = context
-        self.logger = logging.getLogger(self.__class__.__name__)
+        self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
 
     def register_tasks(self, main: Task):
         """
@@ -64,7 +65,7 @@ class ProgressReporter:
         task.success = result.success
         task.summery = result.summery
 
-        report = f"{'\t' * task.depth} finished {task.task_name()} in {task.end_time - task.start_time} with success: {result.success}"
+        report = f"finished {task.task_name()} in {task.end_time - task.start_time} with status: {'success' if result.success else 'failed'}"
         if result.error is not None:
             report += f", error: \n{result.error}"
         else:
@@ -85,7 +86,7 @@ class ProgressReporter:
             task: Task reporting updates.
             batches: Number of batches processed so far.
             expected_batches: Number of expected batches. Can be `None` if the overall number of
-                batches is not know before execution.
+                batches is not known before execution.
             stats: dict of statistics so far (such as `nodes_created`).
         """
         pass
@@ -119,13 +120,16 @@ class Neo4jProgressReporter(ProgressReporter):
             database: Name of the database to write the status updates to.
         """
         super().__init__(context)
+        self.run_uuid = None
         self.database = database
         self.logger.info(f"progress reporting to database: {self.database}")
         self.__create_constraints()
+        self._register_shutdown_handler()
 
     def register_tasks(self, root: Task, **kwargs):
         super().register_tasks(root)
 
+        self.run_uuid = root.uuid
         with self.context.neo4j.session(self.database) as session:
             order = 0
             session.run(
@@ -166,7 +170,7 @@ class Neo4jProgressReporter(ProgressReporter):
                         start_time=task.start_time)
         return task
 
-    def finished_task(self, task: Task,  result: TaskReturn) -> Task:
+    def finished_task(self, task: Task, result: TaskReturn) -> Task:
         super().finished_task(task=task, result=result)
         if result.success:
             status = "success"
@@ -190,6 +194,21 @@ class Neo4jProgressReporter(ProgressReporter):
             session.run("MATCH (t:ETLTask {uuid:$id}) SET t.batches =$batches, t.expected_batches =$expected_batches",
                         id=task.uuid, batches=batches, expected_batches=expected_batches)
 
+    def _register_shutdown_handler(self):
+        def shutdown_handler(signum, frame):
+            self.logger.warning("SIGINT received, waiting for running tasks to abort.")
+            with self.context.neo4j.session(self.database) as session:
+                cnt = session.run("""
+                MATCH path=(r:ETLRun {uuid: $runId})-[*]->() 
+                WITH [task in nodes(path) WHERE task:ETLTask AND task.status IN ['open', 'running'] | task] AS tasks
+                UNWIND tasks AS task
+                SET task.status = 'aborted'
+                RETURN count(task) AS cnt
+                """, runId=self.run_uuid
+                ).single()['cnt']
+            self.logger.info(f"marked {cnt} tasks as aborted.")
+        add_sigint_handler(shutdown_handler)
+
 
 def get_reporter(context) -> ProgressReporter:
     """