neo4j-etl-lib 0.2.0__tar.gz → 0.3.1__tar.gz

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.2.0
+Version: 0.3.1
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>
@@ -14,11 +14,11 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Database
 Classifier: Development Status :: 4 - Beta
 License-File: LICENSE
-Requires-Dist: pydantic>=2.10.5; python_version >= '3.8'
-Requires-Dist: neo4j>=5.27.0; python_version >= '3.7'
-Requires-Dist: python-dotenv>=1.0.1; python_version >= '3.8'
-Requires-Dist: tabulate>=0.9.0; python_version >= '3.7'
-Requires-Dist: click>=8.1.8; python_version >= '3.7'
+Requires-Dist: pydantic>=2.10.5; python_version >= '3.10'
+Requires-Dist: neo4j-rust-ext>=5.27.0,<6; python_version >= '3.10'
+Requires-Dist: python-dotenv>=1.0.1; python_version >= '3.10'
+Requires-Dist: tabulate>=0.9.0; python_version >= '3.10'
+Requires-Dist: click>=8.1.8; python_version >= '3.10'
 Requires-Dist: pydantic[email-validator]
 Requires-Dist: pytest>=8.3.0 ; extra == "dev" and ( python_version >= '3.8')
 Requires-Dist: testcontainers[neo4j]==4.9.0 ; extra == "dev" and ( python_version >= '3.9' and python_version < '4.0')
@@ -35,11 +35,13 @@ Requires-Dist: sphinx-autoapi ; extra == "dev"
 Requires-Dist: sqlalchemy ; extra == "dev"
 Requires-Dist: psycopg2-binary ; extra == "dev"
 Requires-Dist: graphdatascience>=1.13 ; extra == "gds" and ( python_version >= '3.9')
+Requires-Dist: nox>=2024.0.0 ; extra == "nox"
 Requires-Dist: sqlalchemy ; extra == "sql"
 Project-URL: Documentation, https://neo-technology-field.github.io/python-etl-lib/index.html
 Project-URL: Home, https://github.com/neo-technology-field/python-etl-lib
 Provides-Extra: dev
 Provides-Extra: gds
+Provides-Extra: nox
 Provides-Extra: sql
 
 # Neo4j ETL Toolbox

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/pyproject.toml

@@ -22,11 +22,11 @@ dynamic = ["version", "description"]
 keywords = ["etl", "graph", "database"]
 
 dependencies = [
-    "pydantic>=2.10.5; python_version >= '3.8'",
-    "neo4j>=5.27.0; python_version >= '3.7'",
-    "python-dotenv>=1.0.1; python_version >= '3.8'",
-    "tabulate>=0.9.0; python_version >= '3.7'",
-    "click>=8.1.8; python_version >= '3.7'",
+    "pydantic>=2.10.5; python_version >= '3.10'",
+    "neo4j-rust-ext>=5.27.0,<6; python_version >= '3.10'",
+    "python-dotenv>=1.0.1; python_version >= '3.10'",
+    "tabulate>=0.9.0; python_version >= '3.10'",
+    "click>=8.1.8; python_version >= '3.10'",
     "pydantic[email_validator]"
 ]
 
@@ -41,6 +41,11 @@ dev = [
 gds = ["graphdatascience>=1.13; python_version >= '3.9'"]
 sql = ["sqlalchemy"]
 
+# Local-only multy-version testing, install via `pip install ".[dev,nox]"`
+nox = [
+    "nox>=2024.0.0"
+]
+
 [project.urls]
 Home = "https://github.com/neo-technology-field/python-etl-lib"
 Documentation = "https://neo-technology-field.github.io/python-etl-lib/index.html"

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/src/etl_lib/__init__.py

@@ -1,4 +1,4 @@
 """
 Building blocks for ETL pipelines.
 """
-__version__ = "0.2.0"
+__version__ = "0.3.1"

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/src/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."""
 

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/src/etl_lib/core/ETLContext.py

@@ -1,19 +1,23 @@
 import logging
-from typing import NamedTuple, Any
+from typing import Any, Dict, List, NamedTuple
+
+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, WRITE_ACCESS, SummaryCounters
+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
@@ -26,14 +30,29 @@ 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:
@@ -51,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):
@@ -123,6 +143,7 @@ class Neo4jContext:
         self.logger.info(
             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.
@@ -147,14 +168,26 @@ if sqlalchemy_available:
                 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_size=pool_size, max_overflow=max_overflow)
+            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.
     """
 
@@ -163,12 +196,12 @@ 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)
@@ -176,7 +209,7 @@ class ETLContext:
         if sql_uri is not None and sqlalchemy_available:
             self.sql = SQLContext(sql_uri)
         if gds_available:
-            self.gds =gds(self.neo4j)
+            self.gds = gds(self.neo4j)
 
     def env(self, key: str) -> Any:
         """
@@ -190,3 +223,4 @@ class ETLContext:
         """
         if key in self.__env_vars:
             return self.__env_vars[key]
+        return None

neo4j_etl_lib-0.3.1/src/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

{neo4j_etl_lib-0.2.0 → neo4j_etl_lib-0.3.1}/src/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):
         """
@@ -44,7 +45,7 @@ class ProgressReporter:
             The task that was provided.
         """
         task.start_time = datetime.now()
-        self.logger.info(f"{'\t' * task.depth}starting {task.task_name()}")
+        self.logger.info(f"{' ' * (4 * task.depth)}starting {task.task_name()}")
         return task
 
     def finished_task(self, task: Task, result: TaskReturn) -> Task:
@@ -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:
     """