First iteration of TTL fix

Author: angelosantos4

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -254,12 +254,18 @@ def generate_ttl_match_query(self, config: TimeToLiveConfiguration) -> Query:
         return Query(str(query_builder), params)
 
     def generate_ttl_query_from_configuration(
-        self, config: TimeToLiveConfiguration
+        self,
+        config: TimeToLiveConfiguration,
+        chunk_size,
+        execute_chunks_in_parallel,
+        retries_per_chunk,
     ) -> Query:
         ttl_match_query = self.generate_ttl_match_query(config)
         operation = (
             DELETE_NODE_QUERY
             if config.graph_object_type == GraphObjectType.NODE
             else DELETE_REL_QUERY
         )
-        return ttl_match_query.feed_batched_query(operation)
+        return ttl_match_query.feed_batched_query(
+            operation, chunk_size, execute_chunks_in_parallel, retries_per_chunk
+        )

nodestream_plugin_neo4j/query.py

@@ -29,14 +29,23 @@ class Query:
     def from_statement(cls, query_statement: str, **parameters: Any) -> "Query":
         return cls(query_statement, parameters)
 
-    def feed_batched_query(self, batched_query: str) -> "Query":
+    def feed_batched_query(
+        self,
+        batched_query: str,
+        chunk_size: int = 1000,
+        execute_chunks_in_parallel: bool = True,
+        retries_per_chunk: int = 3,
+    ) -> "Query":
         """Feed the results of the the query into another query that will be executed in batches."""
         return Query(
             COMMIT_QUERY,
             {
                 "iterate_params": self.parameters,
                 "batched_query": batched_query,
                 "iterable_query": self.query_statement,
+                "execute_chunks_in_parallel": execute_chunks_in_parallel,
+                "chunk_size": chunk_size,
+                "retries_per_chunk": retries_per_chunk,
             },
         )
 

nodestream_plugin_neo4j/query_executor.py

@@ -71,7 +71,12 @@ async def upsert_relationships_in_bulk_of_same_operation(
         )
 
     async def perform_ttl_op(self, config: TimeToLiveConfiguration):
-        query = self.ingest_query_builder.generate_ttl_query_from_configuration(config)
+        query = self.ingest_query_builder.generate_ttl_query_from_configuration(
+            config,
+            chunk_size=self.chunk_size,
+            execute_chunks_in_parallel=self.execute_chunks_in_parallel,
+            retries_per_chunk=self.retries_per_chunk,
+        )
         await self.database_connection.execute(query)
 
     async def execute_hook(self, hook: IngestionHook):

tests/unit/test_ingest_query_builder.py

@@ -31,6 +31,9 @@ def query_builder():
 
 
 GREATEST_DAY = Timestamp(1998, 3, 25, 2, 0, 1)
+DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL = 1000
+DEFAULT_CHUNK_SIZE = True
+DEFAULT_RETRIES_PER_CHUNK = 3
 
 BASIC_NODE_TTL = TimeToLiveConfiguration(
     graph_object_type=GraphObjectType.NODE,
@@ -43,6 +46,9 @@ def query_builder():
         "iterate_params": {"earliest_allowed_time": GREATEST_DAY},
         "batched_query": DELETE_NODE_QUERY,
         "iterable_query": "MATCH (x: TestNodeType) WHERE x.last_ingested_at <= $earliest_allowed_time RETURN id(x) as id",
+        "execute_chunks_in_parallel": DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
+        "chunk_size": DEFAULT_CHUNK_SIZE,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
 )
 
@@ -58,6 +64,9 @@ def query_builder():
         "iterate_params": {"earliest_allowed_time": GREATEST_DAY},
         "batched_query": DELETE_NODE_QUERY,
         "iterable_query": NODE_TTL_WITH_CUSTOM_QUERY.custom_query,
+        "execute_chunks_in_parallel": DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
+        "chunk_size": DEFAULT_CHUNK_SIZE,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
 )
 
@@ -72,6 +81,9 @@ def query_builder():
         "iterate_params": {"earliest_allowed_time": GREATEST_DAY},
         "iterable_query": "MATCH ()-[x: IS_RELATED_TO]->() WHERE x.last_ingested_at <= $earliest_allowed_time RETURN id(x) as id",
         "batched_query": DELETE_REL_QUERY,
+        "execute_chunks_in_parallel": DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
+        "chunk_size": DEFAULT_CHUNK_SIZE,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
 )
 
@@ -87,6 +99,9 @@ def query_builder():
         "iterate_params": {"earliest_allowed_time": GREATEST_DAY},
         "iterable_query": REL_TTL_WITH_CUSTOM_QUERY.custom_query,
         "batched_query": DELETE_REL_QUERY,
+        "execute_chunks_in_parallel": DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
+        "chunk_size": DEFAULT_CHUNK_SIZE,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
 )
 
@@ -103,7 +118,12 @@ def query_builder():
 )
 def test_generates_expected_queries(mocked_utcnow, query_builder, ttl, expected_query):
     mocked_utcnow.return_value = Timestamp(1998, 3, 25, 12, 0, 1)
-    resultant_query = query_builder.generate_ttl_query_from_configuration(ttl)
+    resultant_query = query_builder.generate_ttl_query_from_configuration(
+        ttl,
+        DEFAULT_CHUNK_SIZE,
+        DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
+        DEFAULT_RETRIES_PER_CHUNK,
+    )
     assert_that(resultant_query, equal_to(expected_query))
 
 

tests/unit/test_query_executor.py

@@ -9,6 +9,12 @@
 
 from .matchers import ran_query
 
+TEST_PARAMS = {
+    "execute_chunks_in_parallel": True,
+    "chunk_size": 1000,
+    "retries_per_chunk": 3,
+}
+
 
 @pytest.fixture
 def some_query():
@@ -84,7 +90,7 @@ async def test_perform_ttl_op(query_executor, some_query):
     )
     query_generator.return_value = some_query
     await query_executor.perform_ttl_op(ttl_config)
-    query_generator.assert_called_once_with(ttl_config)
+    query_generator.assert_called_once_with(ttl_config, **TEST_PARAMS)
     assert_that(query_executor, ran_query(some_query))