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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.1.0
+Version: 0.2.0
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>
@@ -19,6 +19,7 @@ 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[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')
 Requires-Dist: pytest-cov ; extra == "dev"
@@ -31,11 +32,15 @@ Requires-Dist: pydata-sphinx-theme ; extra == "dev"
 Requires-Dist: sphinx-autodoc-typehints ; extra == "dev"
 Requires-Dist: sphinxcontrib-napoleon ; extra == "dev"
 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: 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: sql
 
 # Neo4j ETL Toolbox
 
@@ -43,7 +48,13 @@ A Python library of building blocks to assemble etl pipelines.
 
 Complete documentation can be found on https://neo-technology-field.github.io/python-etl-lib/index.html
 
-See https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/gtfs for an example project.
+See https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/gtfs 
+
+or 
+
+https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/musicbrainz
+
+for example projects.
 
 
 The library can be installed via 

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/README.md

@@ -4,7 +4,13 @@ A Python library of building blocks to assemble etl pipelines.
 
 Complete documentation can be found on https://neo-technology-field.github.io/python-etl-lib/index.html
 
-See https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/gtfs for an example project.
+See https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/gtfs 
+
+or 
+
+https://github.com/neo-technology-field/python-etl-lib/tree/main/examples/musicbrainz
+
+for example projects.
 
 
 The library can be installed via 

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

@@ -26,7 +26,8 @@ dependencies = [
     "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'"
+    "click>=8.1.8; python_version >= '3.7'",
+    "pydantic[email_validator]"
 ]
 
 [project.optional-dependencies]
@@ -35,9 +36,10 @@ dev = [
     "testcontainers[neo4j]==4.9.0; python_version >= '3.9' and python_version < '4.0'",
     "pytest-cov", "bumpver", "isort", "pip-tools",
     "sphinx", "sphinx-rtd-theme", "pydata-sphinx-theme", "sphinx-autodoc-typehints",
-    "sphinxcontrib-napoleon", "sphinx-autoapi"
+    "sphinxcontrib-napoleon", "sphinx-autoapi", "sqlalchemy", "psycopg2-binary"
 ]
 gds = ["graphdatascience>=1.13; python_version >= '3.9'"]
+sql = ["sqlalchemy"]
 
 [project.urls]
 Home = "https://github.com/neo-technology-field/python-etl-lib"

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

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

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/cli/run_tools.py

@@ -55,7 +55,7 @@ def __driver(ctx):
     database_name = ctx.obj["database_name"]
     neo4j_password = ctx.obj["neo4j_password"]
     return GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password), database=database_name,
-                                notifications_min_severity="OFF", user_agent="ETL CLI 0.1")
+                                notifications_min_severity="OFF", user_agent="ETL CLI")
 
 
 @click.group()
@@ -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:
 
@@ -165,25 +165,57 @@ def detail(ctx, run_id, details):
             __print_details(driver, run_id)
 
 
+# noinspection PyTypeChecker
 @cli.command()
-@click.option('--run-id', required=False, help='Run ID to delete')
-@click.option('--since', help='Delete runs since a specific date')
-@click.option('--older', help='Delete runs older than a specific date')
+@click.option('--run-id', required=False, type=str, help='Run IDs to delete, works with comma separated list')
+@click.option('--before', type=click.DateTime(formats=["%Y-%m-%d"]), help='Delete runs before a specific date in format YYYY-MM-DD')
+@click.option('--older', help='Delete runs older than x days', type=int)
 @click.pass_context
-def delete(ctx, run_id, since, older):
+def delete(ctx, run_id, before, older):
     """
-    Delete runs based on run ID, date, or age. One and only one of --run-id, --since, or --older must be provided.
+    Delete runs based on run ID, date, or age. One and only one of --run-id, --before, or --older must be provided.
     """
     # Ensure mutual exclusivity
-    options = [run_id, since, older]
+    options = [run_id, before, older]
     if sum(bool(opt) for opt in options) != 1:
-        print("You must specify exactly one of --run-id, --since, or --older.")
+        print("You must specify exactly one of --run-id, --before, or --older.")
         return
 
     if run_id:
-        print(f"Deleting run ID: {run_id}")
-    elif since:
-        print(f"Deleting runs since: {since}")
+        ids = run_id.split(',')
+        delete_runs(ctx, ids)
+    elif before:
+        print(f"Deleting runs before: {before}")
+        with __driver(ctx) as driver:
+            record= driver.execute_query(
+                """MATCH (r:ETLRun) WHERE date(r.startTime) < date($before) 
+                RETURN collect(r.uuid) AS ids
+                """,
+                result_transformer_=neo4j.Result.single,
+                before=before)
+            ids = record[0]
+            delete_runs(ctx, ids)
+
     elif older:
         print(f"Deleting runs older than: {older}")
-    # Implement delete logic here
+        with __driver(ctx) as driver:
+            record = driver.execute_query(
+                """MATCH (r:ETLRun) WHERE date(r.startTime) < (date() - duration({days: $days})) 
+                RETURN collect(r.uuid) AS ids
+                """,
+                result_transformer_=neo4j.Result.single,
+                days=older)
+            ids = record[0]
+            delete_runs(ctx, ids)
+
+
+def delete_runs(ctx, ids):
+    print(f"Deleting run IDs: {ids}")
+    with __driver(ctx) as driver:
+        records, _, _ = driver.execute_query(
+            """
+            MATCH (r:ETLRun)-[*]->(n) WHERE r.uuid IN $ids
+            DETACH DELETE n
+            DETACH DELETE r
+            """, ids=ids, routing_=neo4j.RoutingControl.WRITE)
+    print(f"Deleted run IDs: {ids}  successfully")

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/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

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

@@ -1,9 +1,26 @@
 import logging
 from typing import NamedTuple, Any
 
-from graphdatascience import GraphDataScience
+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
 
+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
 
 
@@ -99,22 +116,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")
@@ -122,6 +123,32 @@ 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.
+
+    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_size=pool_size, max_overflow=max_overflow)
+
 
 class ETLContext:
     """
@@ -145,6 +172,11 @@ class ETLContext:
         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 +186,7 @@ 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]

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/core/ProgressReporter.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 from tabulate import tabulate
 
-from etl_lib.core.Task import Task, TaskGroup
+from etl_lib.core.Task import Task, TaskGroup, TaskReturn
 
 
 class ProgressReporter:
@@ -47,7 +47,7 @@ class ProgressReporter:
         self.logger.info(f"{'\t' * task.depth}starting {task.task_name()}")
         return task
 
-    def finished_task(self, task: Task, success: bool, summery: dict, error: str = None) -> Task:
+    def finished_task(self, task: Task, result: TaskReturn) -> Task:
         """
         Marks the task as finished.
 
@@ -55,23 +55,21 @@ class ProgressReporter:
 
         Args:
             task: Task to be marked as finished.
-            success: True if the task has successfully finished.
-            summery: statistics for this task (such as `nodes_created`)
-            error: If an exception occurred, the exception text should be provided here.
+            result: result of the task execution, such as status and summery information.
 
         Returns:
             Task to be marked as started.
         """
         task.end_time = datetime.now()
-        task.success = success
-        task.summery = summery
+        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: {success}"
-        if error is not None:
-            report += f", error: \n{error}"
+        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:
             # for the logger, remove entries with 0, but keep them in the original for reporting
-            cleaned_summery = {key: value for key, value in summery.items() if value != 0}
+            cleaned_summery = {key: value for key, value in result.summery.items() if value != 0}
             if len(cleaned_summery) > 0:
                 report += f"\n{tabulate([cleaned_summery], headers='keys', tablefmt='psql')}"
         self.logger.info(report)
@@ -87,7 +85,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
@@ -168,9 +166,9 @@ class Neo4jProgressReporter(ProgressReporter):
                         start_time=task.start_time)
         return task
 
-    def finished_task(self, task: Task, success: bool, summery: dict, error: str = None) -> Task:
-        super().finished_task(task=task, success=success, summery=summery, error=error)
-        if success:
+    def finished_task(self, task: Task,  result: TaskReturn) -> Task:
+        super().finished_task(task=task, result=result)
+        if result.success:
             status = "success"
         else:
             status = "failure"
@@ -179,7 +177,7 @@ class Neo4jProgressReporter(ProgressReporter):
             MATCH (t:ETLTask {uuid:$id}) SET t.endTime = $end_time, t.status = $status, t.error = $error
             CREATE (s:ETLStats) SET s=$summery
             CREATE (t)-[:HAS_STATS]->(s)
-            """, id=task.uuid, end_time=task.end_time, summery=summery, status=status, error=error)
+            """, id=task.uuid, end_time=task.end_time, summery=result.summery, status=status, error=result.error)
         return task
 
     def __create_constraints(self):

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/core/Task.py

@@ -46,7 +46,8 @@ class TaskReturn:
 
         # Combine success values and errors
         combined_success = self.success and other.success
-        combined_error = f"{self.error or ''} | {other.error or ''}".strip(" |")
+        combined_error = None if not (self.error or other.error) \
+            else f"{self.error or ''} | {other.error or ''}".strip(" |")
 
         return TaskReturn(
             success=combined_success, summery=merged_summery, error=combined_error
@@ -99,12 +100,7 @@ class Task:
         except Exception as e:
             result = TaskReturn(success=False, summery={}, error=str(e))
 
-        self.context.reporter.finished_task(
-            task=self,
-            success=result.success,
-            summery=result.summery,
-            error=result.error,
-        )
+        self.context.reporter.finished_task(task=self,result=result)
 
         return result
 

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/data_sink/CypherBatchSink.py

@@ -10,7 +10,7 @@ class CypherBatchSink(BatchProcessor):
     BatchProcessor to write batches of data to a Neo4j database.
     """
 
-    def __init__(self, context: ETLContext, task: Task, predecessor: BatchProcessor, query: str):
+    def __init__(self, context: ETLContext, task: Task, predecessor: BatchProcessor, query: str, **kwargs):
         """
         Constructs a new CypherBatchSink.
 
@@ -20,16 +20,17 @@ class CypherBatchSink(BatchProcessor):
             predecessor: BatchProcessor which :func:`~get_batch` function will be called to receive batches to process.
             query: Cypher to write the query to Neo4j.
                 Data will be passed as `batch` parameter.
-                Therefor, the query should start with a `UNWIND $batch AS row`.
+                Therefore, the query should start with a `UNWIND $batch AS row`.
         """
         super().__init__(context, task, predecessor)
         self.query = query
         self.neo4j = context.neo4j
+        self.kwargs = kwargs
 
     def get_batch(self, batch_size: int) -> Generator[BatchResults, None, None]:
         assert self.predecessor is not None
 
         with self.neo4j.session() as session:
             for batch_result in self.predecessor.get_batch(batch_size):
-                result = self.neo4j.query_database(session=session, query=self.query, batch=batch_result.chunk)
+                result = self.neo4j.query_database(session=session, query=self.query, batch=batch_result.chunk, **self.kwargs)
                 yield append_result(batch_result, result.summery)

neo4j_etl_lib-0.2.0/src/etl_lib/data_sink/SQLBatchSink.py

@@ -0,0 +1,36 @@
+from typing import Generator
+from sqlalchemy import text
+from etl_lib.core.ETLContext import ETLContext
+from etl_lib.core.BatchProcessor import BatchProcessor, BatchResults, append_result
+from etl_lib.core.Task import Task
+
+
+class SQLBatchSink(BatchProcessor):
+    """
+    BatchProcessor to write batches of data to an SQL database.
+    """
+
+    def __init__(self, context: ETLContext, task: Task, predecessor: BatchProcessor, query: str):
+        """
+        Constructs a new SQLBatchSink.
+
+        Args:
+            context: ETLContext instance.
+            task: Task instance owning this batchProcessor.
+            predecessor: BatchProcessor which `get_batch` function will be called to receive batches to process.
+            query: SQL query to write data.
+                Data will be passed as a batch using parameterized statements (`:param_name` syntax).
+        """
+        super().__init__(context, task, predecessor)
+        self.query = query
+        self.engine = context.sql.engine
+
+    def get_batch(self, batch_size: int) -> Generator[BatchResults, None, None]:
+        assert self.predecessor is not None
+
+        with self.engine.connect() as conn:
+            with conn.begin():
+                for batch_result in self.predecessor.get_batch(batch_size):
+                    conn.execute(text(self.query), batch_result.chunk)
+                    yield append_result(batch_result, {"sql_rows_written": len(batch_result.chunk)})
+

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/data_source/CypherBatchSource.py

@@ -1,4 +1,6 @@
-from typing import Generator
+from typing import Generator, Callable, Optional
+
+from neo4j import Record
 
 from etl_lib.core.BatchProcessor import BatchResults, BatchProcessor
 from etl_lib.core.ETLContext import ETLContext
@@ -7,7 +9,14 @@ from etl_lib.core.Task import Task
 
 class CypherBatchSource(BatchProcessor):
 
-    def __init__(self, context: ETLContext, task: Task, query: str, **kwargs):
+    def __init__(
+            self,
+            context: ETLContext,
+            task: Task,
+            query: str,
+            record_transformer: Optional[Callable[[Record], dict]] = None,
+            **kwargs
+    ):
         """
        Constructs a new CypherBatchSource.
 
@@ -15,10 +24,12 @@ class CypherBatchSource(BatchProcessor):
            context: :class:`etl_lib.core.ETLContext.ETLContext` instance.
            task: :class:`etl_lib.core.Task.Task` instance owning this batchProcessor.
            query: Cypher query to execute.
+           record_transformer: Optional function to transform each record. See Neo4j API documentation on `result_transformer_`
            kwargs: Arguments passed as parameters with the query.
        """
         super().__init__(context, task)
         self.query = query
+        self.record_transformer = record_transformer
         self.kwargs = kwargs
 
     def __read_records(self, tx, batch_size):
@@ -26,7 +37,11 @@ class CypherBatchSource(BatchProcessor):
         result = tx.run(self.query, **self.kwargs)
 
         for record in result:
-            batch_.append(record.data())
+            data = record.data()
+            if self.record_transformer:
+                data = self.record_transformer(data)
+            batch_.append(data)
+
             if len(batch_) == batch_size:
                 yield batch_
                 batch_ = []

neo4j_etl_lib-0.2.0/src/etl_lib/data_source/SQLBatchSource.py

@@ -0,0 +1,60 @@
+from typing import Generator, Callable, Optional
+from sqlalchemy import text
+from etl_lib.core.BatchProcessor import BatchResults, BatchProcessor
+from etl_lib.core.ETLContext import ETLContext
+from etl_lib.core.Task import Task
+
+
+class SQLBatchSource(BatchProcessor):
+    def __init__(
+            self,
+            context: ETLContext,
+            task: Task,
+            query: str,
+            record_transformer: Optional[Callable[[dict], dict]] = None,
+            **kwargs
+    ):
+        """
+        Constructs a new SQLBatchSource.
+
+        Args:
+            context: :class:`etl_lib.core.ETLContext.ETLContext` instance.
+            task: :class:`etl_lib.core.Task.Task` instance owning this batchProcessor.
+            query: SQL query to execute.
+            record_transformer: Optional function to transform each row (dict format).
+            kwargs: Arguments passed as parameters with the query.
+        """
+        super().__init__(context, task)
+        self.query = query
+        self.record_transformer = record_transformer
+        self.kwargs = kwargs  # Query parameters
+
+    def __read_records(self, conn, batch_size: int):
+        batch_ = []
+        result = conn.execute(text(self.query), self.kwargs)  # Safe execution with bound parameters
+
+        for row in result.mappings():  # Returns row as dict (like Neo4j's `record.data()`)
+            data = dict(row)  # Convert to dictionary
+            if self.record_transformer:
+                data = self.record_transformer(data)
+            batch_.append(data)
+
+            if len(batch_) == batch_size:
+                yield batch_
+                batch_ = []  # Reset batch
+
+        if batch_:
+            yield batch_
+
+    def get_batch(self, max_batch_size: int) -> Generator[BatchResults, None, None]:
+        """
+        Fetches data in batches using an open transaction, similar to Neo4j's approach.
+        """
+        with self.context.sql.engine.connect() as conn:  # Keep transaction open
+            with conn.begin():  # Ensures rollback on failure
+                for chunk in self.__read_records(conn, max_batch_size):
+                    yield BatchResults(
+                        chunk=chunk,
+                        statistics={"sql_rows_read": len(chunk)},
+                        batch_size=len(chunk)
+                    )

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/task/CreateReportingConstraintsTask.py

@@ -4,8 +4,8 @@ from etl_lib.core.Task import Task, TaskReturn
 class CreateReportingConstraintsTask(Task):
     """Creates the constraint in the REPORTER_DATABASE database."""
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, context):
+        super().__init__(context)
 
     def run_internal(self, **kwargs) -> TaskReturn:
         database = self.context.env("REPORTER_DATABASE")

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/task/ExecuteCypherTask.py

@@ -24,10 +24,10 @@ class ExecuteCypherTask(Task):
                 for query in self._query():
                     result = self.context.neo4j.query_database(session=session, query=query, **kwargs)
                     stats = merge_summery(stats, result.summery)
-                return TaskReturn(True, stats)
+                return TaskReturn(success=True, summery=stats)
             else:
                 result = self.context.neo4j.query_database(session=session, query=self._query(), **kwargs)
-                return TaskReturn(True, result.summery)
+                return TaskReturn(success=True, summery=result.summery)
 
     @abc.abstractmethod
     def _query(self) -> str | list[str]:

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/task/data_loading/CSVLoad2Neo4jTask.py

@@ -14,19 +14,55 @@ from etl_lib.data_source.CSVBatchSource import CSVBatchSource
 
 
 class CSVLoad2Neo4jTask(Task):
-    """
+    '''
     Loads the specified CSV file to Neo4j.
 
     Uses BatchProcessors to read, validate and write to Neo4j.
     The validation step is using pydantic, hence a Pydantic model needs to be provided.
-    Rows that fail the validation, will be written to en error file. The location of the error file is determined as
+    Rows with fail validation will be written to en error file. The location of the error file is determined as
     follows:
 
-    If the context env vars hold an entry `ETL_ERROR_PATH` the file will be place there, with the name set to name
+    If the context env vars hold an entry `ETL_ERROR_PATH` the file will be placed there, with the name set to name
     of the provided filename appended with `.error.json`
 
-    If  `ETL_ERROR_PATH` is not set, the file will be placed in the same directory as the CSV file.
-    """
+    If `ETL_ERROR_PATH` is not set, the file will be placed in the same directory as the CSV file.
+
+    Example usage: (from the gtfs demo)
+
+    .. code-block:: python
+
+        class LoadStopsTask(CSVLoad2Neo4jTask):
+            class Stop(BaseModel):
+                id: str = Field(alias="stop_id")
+                name: str = Field(alias="stop_name")
+                latitude: float = Field(alias="stop_lat")
+                longitude: float = Field(alias="stop_lon")
+                platform_code: Optional[str] = None
+                parent_station: Optional[str] = None
+                type: Optional[str] = Field(alias="location_type", default=None)
+                timezone: Optional[str] = Field(alias="stop_timezone", default=None)
+                code: Optional[str] = Field(alias="stop_code", default=None)
+
+            def __init__(self, context: ETLContext, file: Path):
+                super().__init__(context, LoadStopsTask.Stop, file)
+
+            def task_name(self) -> str:
+                return f"{self.__class__.__name__}('{self.file}')"
+
+            def _query(self):
+                return """
+                       UNWIND $batch AS row
+                       MERGE (s:Stop {id: row.id})
+                         SET s.name = row.name,
+                             s.location= point({latitude: row.latitude, longitude: row.longitude}),
+                             s.platformCode= row.platform_code,
+                             s.parentStation= row.parent_station,
+                             s.type= row.type,
+                             s.timezone= row.timezone,
+                             s.code= row.code
+                      """
+
+    '''
     def __init__(self, context: ETLContext, model: Type[BaseModel], file: Path, batch_size: int = 5000):
         super().__init__(context)
         self.batch_size = batch_size

neo4j_etl_lib-0.2.0/src/etl_lib/task/data_loading/SQLLoad2Neo4jTask.py

@@ -0,0 +1,90 @@
+from abc import abstractmethod
+
+from sqlalchemy import text
+
+from etl_lib.core import ETLContext
+from etl_lib.core.ClosedLoopBatchProcessor import ClosedLoopBatchProcessor
+from etl_lib.core.Task import Task, TaskReturn
+from etl_lib.data_sink.CypherBatchSink import CypherBatchSink
+from etl_lib.data_source.SQLBatchSource import SQLBatchSource
+
+
+class SQLLoad2Neo4jTask(Task):
+    '''
+    Load the output of the specified SQL query to Neo4j.
+
+    Uses BatchProcessors to read and write data.
+    Subclasses must implement the methods returning the SQL and Cypher queries.
+
+    Example usage: (from the MusicBrainz example)
+
+    .. code-block:: python
+
+
+        class LoadArtistCreditTask(SQLLoad2Neo4jTask):
+            def _sql_query(self) -> str:
+                return """
+                    SELECT ac.id AS artist_credit_id, ac.name AS credit_name
+                    FROM artist_credit ac;
+                    """
+
+            def _cypher_query(self) -> str:
+                return """
+                       UNWIND $batch AS row
+                       MERGE (ac:ArtistCredit {id: row.artist_credit_id})
+                       SET ac.name = row.credit_name
+                      """
+
+            def _count_query(self) -> str | None:
+                return "SELECT COUNT(*) FROM artist_credit;"
+
+    '''
+
+    def __init__(self, context: ETLContext, batch_size: int = 5000):
+        super().__init__(context)
+        self.context = context
+        self.batch_size = batch_size
+
+    @abstractmethod
+    def _sql_query(self) -> str:
+        """
+        Return the SQL query to load the source data.
+        """
+        pass
+
+    @abstractmethod
+    def _cypher_query(self) -> str:
+        """
+        Return the Cypher query to write the data in batches to Neo4j.
+        """
+        pass
+
+    def _count_query(self) -> str | None:
+        """
+        Return the SQL query to count the number of rows returned from :func:`_sql_query`.
+
+        Optional. If provided, it will run once at the beginning of the task and
+        provide the :class:`etl_lib.core.ClosedLoopBatchProcessor` with the total number of rows.
+        """
+        return None
+
+    def run_internal(self) -> TaskReturn:
+        total_count = self.__get_source_count()
+        source = SQLBatchSource(self.context, self, self._sql_query())
+        sink = CypherBatchSink(self.context, self, source, self._cypher_query())
+
+        end = ClosedLoopBatchProcessor(self.context, self, sink, total_count)
+
+        result = next(end.get_batch(self.batch_size))
+        return TaskReturn(True, result.statistics)
+
+    def __get_source_count(self):
+        count_query = self._count_query()
+        if count_query is None:
+            return None
+
+        with self.context.sql.engine.connect() as conn:
+            with conn.begin():
+                result = conn.execute(text(count_query))
+                row = result.fetchone()
+                return row[0] if row else None

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.2.0}/src/etl_lib/test_utils/utils.py

@@ -7,7 +7,7 @@ from _pytest.tmpdir import tmp_path
 from neo4j import Driver
 from neo4j.time import Date
 
-from etl_lib.core.ETLContext import QueryResult, Neo4jContext, ETLContext
+from etl_lib.core.ETLContext import QueryResult, Neo4jContext, ETLContext, SQLContext, gds
 from etl_lib.core.Task import Task
 
 
@@ -102,6 +102,7 @@ class TestNeo4jContext(Neo4jContext):
         self.logger = logging.getLogger(self.__class__.__name__)
         self.driver = driver
         self.database = get_database_name()
+        self.gds = gds(self)
 
 
 class TestETLContext(ETLContext):
@@ -116,6 +117,16 @@ class TestETLContext(ETLContext):
         if key in self.__env_vars:
             return self.__env_vars[key]
 
+class TestSQLETLContext(ETLContext):
+
+    def __init__(self, sql_uri):
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.reporter = DummyReporter()
+        self.sql = SQLContext(sql_uri)
+
+    def env(self, key: str) -> Any:
+        if key in self.__env_vars:
+            return self.__env_vars[key]
 
 class DummyReporter:
 
@@ -125,7 +136,7 @@ class DummyReporter:
     def started_task(self, task: Task) -> Task:
         pass
 
-    def finished_task(self, task, success: bool, summery: dict, error: str = None) -> Task:
+    def finished_task(self, task, result) -> Task:
         pass
 
     def report_progress(self, task, batches: int, expected_batches: int, stats: dict) -> None:
@@ -151,3 +162,10 @@ class DummyContext:
 
     def env(self, key: str) -> Any:
         pass
+
+class DummyPredecessor:
+    def __init__(self, batches):
+        self.batches = batches
+
+    def get_batch(self, batch_size):
+        yield from self.batches