neo4j-etl-lib 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

etl_lib/task/data_loading/ParallelSQLLoad2Neo4jTask.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Union
+
+from etl_lib.core.ClosedLoopBatchProcessor import ClosedLoopBatchProcessor
+from etl_lib.core.ETLContext import ETLContext
+from etl_lib.core.ParallelBatchProcessor import ParallelBatchProcessor
+from etl_lib.core.SplittingBatchProcessor import SplittingBatchProcessor, dict_id_extractor
+from etl_lib.core.Task import Task, TaskReturn
+from etl_lib.data_sink.CypherBatchSink import CypherBatchSink
+from etl_lib.data_source.SQLBatchSource import SQLBatchSource
+from sqlalchemy import text
+
+
+class ParallelSQLLoad2Neo4jTask(Task, ABC):
+    """
+    Parallelized version of SQLLoad2Neo4jTask: reads via SQLBatchSource,
+    splits into non-overlapping partitions (grid), processes each partition
+    in parallel through a CypherBatchSink, and closes the loop.
+
+    Subclasses must implement:
+        - _sql_query()
+        - _cypher_query()
+        - optionally override _count_query() and _id_extractor().
+
+    Control parameters:
+        batch_size: max items per partition batch
+        table_size: dimension of the splitting grid
+        max_workers: parallel threads per partition group (defaults to table_size)
+        prefetch: number of partition-groups to prefetch
+    """
+
+    def __init__(
+            self,
+            context: ETLContext,
+            batch_size: int = 5000,
+            table_size: int = 10,
+            max_workers: int = None,
+            prefetch: int = 4
+    ):
+        super().__init__(context)
+        self.context = context
+        self.batch_size = batch_size
+        self.table_size = table_size
+        # default max_workers to table_size for full parallelism
+        self.max_workers = max_workers or table_size
+        self.prefetch = prefetch
+
+    @abstractmethod
+    def _sql_query(self) -> str:
+        """
+        Return the SQL query to load source rows.
+        """
+        pass
+
+    @abstractmethod
+    def _cypher_query(self) -> str:
+        """
+        Return the Cypher query to write rows into Neo4j.
+        """
+        pass
+
+    def _count_query(self) -> Union[str, None]:
+        """
+        Optional SQL to count source rows for progress reporting.
+        """
+        return None
+
+    def _id_extractor(self) -> Callable:
+        """
+        Extractor mapping each row item to a (row,col) partition index.
+        Default expects dict rows with 'start' and 'end' keys.
+        Override to customize.
+        """
+        return dict_id_extractor()
+
+    def run_internal(self) -> TaskReturn:
+        # total count for ClosedLoopBatchProcessor
+        total_count = self.__get_source_count()
+        # source of raw rows
+        source = SQLBatchSource(self.context, self, self._sql_query())
+
+        # splitter: non-overlapping partitions as defined by the id_extractor
+        splitter = SplittingBatchProcessor(
+            context=self.context,
+            task=self,
+            predecessor=source,
+            table_size=self.table_size,
+            id_extractor=self._id_extractor()
+        )
+
+        # parallel processor: runs CypherBatchSink on each partition concurrently
+        parallel = ParallelBatchProcessor(
+            context=self.context,
+            task=self,
+            worker_factory=lambda: CypherBatchSink(context=self.context, task=self, predecessor=None,
+                                                   query=self._cypher_query()),
+            predecessor=splitter,
+            max_workers=self.max_workers,
+            prefetch=self.prefetch
+        )
+
+        # close loop: drives the pipeline and reports progress
+        closing = ClosedLoopBatchProcessor(
+            context=self.context,
+            task=self,
+            predecessor=parallel,
+            expected_rows=total_count
+        )
+
+        # run once to completion and return aggregated stats
+        result = next(closing.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

etl_lib/task/data_loading/SQLLoad2Neo4jTask.py

@@ -1,8 +1,8 @@
 from abc import abstractmethod
 
+from etl_lib.core.ETLContext import ETLContext
 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

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, SQLContext, gds
+from etl_lib.core.ETLContext import ETLContext, Neo4jContext, QueryResult, SQLContext, gds
 from etl_lib.core.Task import Task
 
 
@@ -96,7 +96,7 @@ def get_database_name():
             raise Exception("define NEO4J_TEST_DATABASE environment variable")
 
 
-class TestNeo4jContext(Neo4jContext):
+class MockNeo4jContext(Neo4jContext):
 
     def __init__(self, driver: Driver):
         self.logger = logging.getLogger(self.__class__.__name__)
@@ -105,29 +105,32 @@ class TestNeo4jContext(Neo4jContext):
         self.gds = gds(self)
 
 
-class TestETLContext(ETLContext):
+class MockETLContext(ETLContext):
 
     def __init__(self, driver: Driver, tmp_path):
         self.logger = logging.getLogger(self.__class__.__name__)
         self.__env_vars = {"ETL_ERROR_PATH": tmp_path}
-        self.neo4j = TestNeo4jContext(driver)
+        self.neo4j = MockNeo4jContext(driver)
         self.reporter = DummyReporter()
 
     def env(self, key: str) -> Any:
         if key in self.__env_vars:
             return self.__env_vars[key]
 
-class TestSQLETLContext(ETLContext):
+
+class MockSQLETLContext(ETLContext):
 
     def __init__(self, sql_uri):
         self.logger = logging.getLogger(self.__class__.__name__)
         self.reporter = DummyReporter()
+        self.__env_vars = {}
         self.sql = SQLContext(sql_uri)
 
     def env(self, key: str) -> Any:
         if key in self.__env_vars:
             return self.__env_vars[key]
 
+
 class DummyReporter:
 
     def register_tasks(self, main: Task):
@@ -163,6 +166,7 @@ class DummyContext:
     def env(self, key: str) -> Any:
         pass
 
+
 class DummyPredecessor:
     def __init__(self, batches):
         self.batches = batches

{neo4j_etl_lib-0.2.0.dist-info → neo4j_etl_lib-0.3.1.dist-info}/METADATA

@@ -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.3.1.dist-info/RECORD

@@ -0,0 +1,36 @@
+etl_lib/__init__.py,sha256=x6coFV38ytJ_wPhR3c0UEzX65oTz2ouKwygkC_tyRLM,65
+etl_lib/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/cli/run_tools.py,sha256=KIar-y22P4kKm-yoJjecYsPwqC7U76M71dEgFO5-ZBo,8561
+etl_lib/core/BatchProcessor.py,sha256=mRpdxZ6ZMKI8XsY3TPuy4dVcvRqLKCO-p63KeOhFyKE,3417
+etl_lib/core/ClosedLoopBatchProcessor.py,sha256=WzML1nldhZRbP8fhlD6utuK5SBYRl1cJgEobVDIdBP4,1626
+etl_lib/core/ETLContext.py,sha256=wmEnbs3n_80B6La9Py_-MHG8BN0FajE9MjGPej0A3To,8045
+etl_lib/core/ParallelBatchProcessor.py,sha256=jNo1Xv1Ts34UZIseoQLDZOhHOVeEr8dUibKUt0FJ4Hw,7318
+etl_lib/core/ProgressReporter.py,sha256=tkE-W6qlR25nU8nUoECcxZDnjnG8AtQH9s9s5WBh_-Q,9377
+etl_lib/core/SplittingBatchProcessor.py,sha256=OIRMUVFpUoZc0w__JJjUr7B9QC3sBlqQp41xghrQzC0,11616
+etl_lib/core/Task.py,sha256=muQFY5qj2n-ZVV8F6vlHqo2lVSvB3wtGdIgkSXVpOFM,9365
+etl_lib/core/ValidationBatchProcessor.py,sha256=U1M2Qp9Ledt8qFiHAg8zMxE9lLRkBrr51NKs_Y8skK8,3400
+etl_lib/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/core/utils.py,sha256=9gezAateAmtWZ5lr6ZWYlTtgrj4-CUKJSGz7g1upnRg,2293
+etl_lib/data_sink/CSVBatchSink.py,sha256=oq4VJwnA4WSyJzdvwstGv73vOEuWmPSfCynhVmxBByU,2204
+etl_lib/data_sink/CypherBatchSink.py,sha256=nBH4bzN1IvdSFcKgiAIrAY5IauB565sdyVrnRc1hg_4,1566
+etl_lib/data_sink/SQLBatchSink.py,sha256=vyGrrxpdmCLUZMI2_W2ORej3FLGbwN9-b2GMYHd-k9g,1451
+etl_lib/data_sink/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/data_source/CSVBatchSource.py,sha256=0q1XdPhAIKw1HcTpnp_F4WxRUzk-24Q8Qd-WeIo5OZ0,2779
+etl_lib/data_source/CypherBatchSource.py,sha256=06WuW11BqYjAXBZqL96Qr9MR8JrcjujDpxXe8cI-SYY,2238
+etl_lib/data_source/SQLBatchSource.py,sha256=O3ZA2GXvo5j_KGwOILzguYZMPY_FJkV5j8FIa3-d9oM,4067
+etl_lib/data_source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/task/CreateReportingConstraintsTask.py,sha256=nTcHLBIgXz_h2OQg-SHjQr68bhH974u0MwrtWPnVwng,762
+etl_lib/task/ExecuteCypherTask.py,sha256=thE8YTZzv1abxNhhDcb4p4ke6qmI6kWR4XQ-GrCBBBU,1284
+etl_lib/task/GDSTask.py,sha256=UP_NMvdeQ9ueLUFlREfe0q3XhFHTCaMvXArSpvxZNiQ,1918
+etl_lib/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/task/data_loading/CSVLoad2Neo4jTask.py,sha256=9XiVdJHpABE-Hx1bsvTKLJWtChc8XMwXeO5RicaHDUo,3873
+etl_lib/task/data_loading/ParallelCSVLoad2Neo4jTask.py,sha256=2xN-5bHV9XgoaJLbbTEBuJFoZHV_CYi_hg6M1HQ-ffA,4030
+etl_lib/task/data_loading/ParallelSQLLoad2Neo4jTask.py,sha256=9xpCW5i8yGnUHyg475Z8rjN2P5NDshJGUEgTU6sm2Bs,4286
+etl_lib/task/data_loading/SQLLoad2Neo4jTask.py,sha256=HR3DcjOUkQN4SbCkgQYzljQCYhOcb-x2-DR5dBdapzU,2953
+etl_lib/task/data_loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+etl_lib/test_utils/utils.py,sha256=CgYOCXcUyndOdRAmGyPLoCIuEik0yzy6FLV2k16cpDM,5673
+neo4j_etl_lib-0.3.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+neo4j_etl_lib-0.3.1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+neo4j_etl_lib-0.3.1.dist-info/METADATA,sha256=Pm921qyxL36Ed_Ppp2cW3OFPxUGMv7IyRTmtba3n96o,2580
+neo4j_etl_lib-0.3.1.dist-info/RECORD,,

neo4j_etl_lib-0.2.0.dist-info/RECORD

@@ -1,32 +0,0 @@
-etl_lib/__init__.py,sha256=m8CiP9nDQUZCy4l5Crx09vtYrpAtqWkstIOj_Wlp1K8,65
-etl_lib/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/cli/run_tools.py,sha256=KIar-y22P4kKm-yoJjecYsPwqC7U76M71dEgFO5-ZBo,8561
-etl_lib/core/BatchProcessor.py,sha256=6quNPE9Dp8hYJDQDTqxQtxbQ3KCmb56Mko34EIsNhyI,3352
-etl_lib/core/ClosedLoopBatchProcessor.py,sha256=WzML1nldhZRbP8fhlD6utuK5SBYRl1cJgEobVDIdBP4,1626
-etl_lib/core/ETLContext.py,sha256=CHRNfjTmIIrPUp2F7l4MJlTsRLwfuBsNJdF_SEeShvI,6695
-etl_lib/core/ProgressReporter.py,sha256=FNmunA8q1vkAcmDF5z8ExH-mTotQMX0fyf2VAz-zI8o,8414
-etl_lib/core/Task.py,sha256=3e8iVXSfXaeBecvgTcs2LiIf2JwpKETRFhH4ig6lock,9202
-etl_lib/core/ValidationBatchProcessor.py,sha256=EhO6PFQB-4PZgIOTXP4PwkbAl5HRK0zgTeKMseRU5QU,3261
-etl_lib/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/core/utils.py,sha256=wwfyvy78fL6sqHdV0IFqAVyEkp6vo5Yo8gRZua2dulw,816
-etl_lib/data_sink/CSVBatchSink.py,sha256=oq4VJwnA4WSyJzdvwstGv73vOEuWmPSfCynhVmxBByU,2204
-etl_lib/data_sink/CypherBatchSink.py,sha256=nBH4bzN1IvdSFcKgiAIrAY5IauB565sdyVrnRc1hg_4,1566
-etl_lib/data_sink/SQLBatchSink.py,sha256=vyGrrxpdmCLUZMI2_W2ORej3FLGbwN9-b2GMYHd-k9g,1451
-etl_lib/data_sink/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/data_source/CSVBatchSource.py,sha256=HILkaQAFua1OM7xgSNKY6teXpcJjWUPaS4Aol-GLYL8,2767
-etl_lib/data_source/CypherBatchSource.py,sha256=06WuW11BqYjAXBZqL96Qr9MR8JrcjujDpxXe8cI-SYY,2238
-etl_lib/data_source/SQLBatchSource.py,sha256=rfIpzMU9IG5xrftZss7W-KpSIINCUdPfJBEXa5QDFts,2323
-etl_lib/data_source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/task/CreateReportingConstraintsTask.py,sha256=nTcHLBIgXz_h2OQg-SHjQr68bhH974u0MwrtWPnVwng,762
-etl_lib/task/ExecuteCypherTask.py,sha256=thE8YTZzv1abxNhhDcb4p4ke6qmI6kWR4XQ-GrCBBBU,1284
-etl_lib/task/GDSTask.py,sha256=X1E83wYa-N7AXy43WPEqIy77d__z-2wpBjWNhGNXJzA,1781
-etl_lib/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/task/data_loading/CSVLoad2Neo4jTask.py,sha256=bD1dWX3guY0f_lFqQP_pzPa03AYveLhjDQzOqkyJVQc,3840
-etl_lib/task/data_loading/SQLLoad2Neo4jTask.py,sha256=N8pQFx5vYKJ3UDn6sw7yBmCJao_b5vBOT3tJ8pWVVLk,2942
-etl_lib/task/data_loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/test_utils/utils.py,sha256=kxWJqdRf1pg-4ByMfrtW3HDbgXIvyVtLndGDVvMCmoI,5641
-neo4j_etl_lib-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-neo4j_etl_lib-0.2.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-neo4j_etl_lib-0.2.0.dist-info/METADATA,sha256=FISyLgs53iHj3Uun7RgGFVtltLus_FiCkU591q9fJ3o,2497
-neo4j_etl_lib-0.2.0.dist-info/RECORD,,