Added in the tests

Author: angelosantos4

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -256,16 +256,17 @@ def generate_ttl_match_query(self, config: TimeToLiveConfiguration) -> Query:
     def generate_ttl_query_from_configuration(
         self,
         config: TimeToLiveConfiguration,
-        chunk_size,
-        execute_chunks_in_parallel,
         retries_per_chunk,
     ) -> Query:
         ttl_match_query = self.generate_ttl_match_query(config)
+        execute_chunks_in_parallel = (
+            config.graph_object_type == GraphObjectType.RELATIONSHIP
+        )
         operation = (
             DELETE_NODE_QUERY
             if config.graph_object_type == GraphObjectType.NODE
             else DELETE_REL_QUERY
         )
         return ttl_match_query.feed_batched_query(
-            operation, chunk_size, execute_chunks_in_parallel, retries_per_chunk
+            operation, config.batch_size, execute_chunks_in_parallel, retries_per_chunk
         )

nodestream_plugin_neo4j/query_executor.py

@@ -73,8 +73,6 @@ 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,
-            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)

tests/e2e/test_neo4j_pipelines.py

@@ -177,6 +177,32 @@ def create_relationship_query_from_params(
     )
 
 
+"""
+The test works like this:
+    We create 60 nodes in total:
+        10 ObjectA 60hrs old
+        10 ObjectB 60hrs old
+        10 ObjectA 36hrs old
+        10 ObjectB 36hrs old
+        10 ObjectA 12hrs old
+        10 ObjectB 12hrs old
+
+    We create 60 pairwise relationships as follows:
+        ObjectA -ADJACENT_TO {60hrs old} > Object B
+        ObjectA -CONNECTED_TO {60hrs old} > Object B
+        ....
+
+    Following the same pattern.
+
+    With our TTL configuration, were able to first get rid of relationships.
+    CONNECTED_TO has 48hr expiry leaving 20 relationships of this type post deletion.
+    ADJACENT_TO has 24hr expiry leaving 10 relationships of this type post deletion
+
+    We assert that the nodes were untouched in the validate_consistency_in_node_counts  function located above.
+    We then do the exact same thing with the nodes.
+"""
+
+
 def create_test_objects(session: Session):
     def create_node(node_label, node_id, timestamp):
         query = create_node_query_from_params(node_label, node_id, timestamp)

tests/unit/test_ingest_query_builder.py

@@ -31,8 +31,7 @@ def query_builder():
 
 
 GREATEST_DAY = Timestamp(1998, 3, 25, 2, 0, 1)
-DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL = True
-DEFAULT_CHUNK_SIZE = 1000
+DEFAULT_CHUNK_SIZE = 100
 DEFAULT_RETRIES_PER_CHUNK = 3
 
 BASIC_NODE_TTL = TimeToLiveConfiguration(
@@ -46,7 +45,7 @@ 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,
+        "execute_chunks_in_parallel": False,
         "chunk_size": DEFAULT_CHUNK_SIZE,
         "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
@@ -64,7 +63,7 @@ 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,
+        "execute_chunks_in_parallel": False,
         "chunk_size": DEFAULT_CHUNK_SIZE,
         "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
@@ -81,7 +80,7 @@ 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,
+        "execute_chunks_in_parallel": True,
         "chunk_size": DEFAULT_CHUNK_SIZE,
         "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
@@ -99,12 +98,50 @@ 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,
+        "execute_chunks_in_parallel": True,
         "chunk_size": DEFAULT_CHUNK_SIZE,
         "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
     },
 )
 
+NODE_TTL_WITH_BATCH_PARAMS = TimeToLiveConfiguration(
+    graph_object_type=GraphObjectType.NODE,
+    object_type="TestNodeType",
+    expiry_in_hours=10,
+    batch_size=200,
+)
+
+NODE_TTL_WITH_BATCH_PARAMS_EXPECTED_QUERY = Query(
+    COMMIT_QUERY,
+    {
+        "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": False,
+        "chunk_size": 200,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
+    },
+)
+
+REL_TTL_WITH_BATCH_PARAMS = TimeToLiveConfiguration(
+    graph_object_type=GraphObjectType.RELATIONSHIP,
+    object_type="IS_RELATED_TO",
+    expiry_in_hours=10,
+    batch_size=200,
+)
+
+REL_TTL_WITH_BATCH_PARAMS_EXPECTED_QUERY = Query(
+    COMMIT_QUERY,
+    {
+        "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": True,
+        "chunk_size": 200,
+        "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
+    },
+)
+
 
 @patch("pandas.Timestamp.utcnow")
 @pytest.mark.parametrize(
@@ -114,14 +151,14 @@ def query_builder():
         (NODE_TTL_WITH_CUSTOM_QUERY, NODE_TTL_WITH_CUSTOM_QUERY_EXPECTED_QUERY),
         (BASIC_REL_TTL, BASIC_REL_TTL_EXPECTED_QUERY),
         (REL_TTL_WITH_CUSTOM_QUERY, REL_TTL_WITH_CUSTOM_QUERY_EXPECTED_QUERY),
+        (NODE_TTL_WITH_BATCH_PARAMS, NODE_TTL_WITH_BATCH_PARAMS_EXPECTED_QUERY),
+        (REL_TTL_WITH_BATCH_PARAMS, REL_TTL_WITH_BATCH_PARAMS_EXPECTED_QUERY),
     ],
 )
 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,
-        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

@@ -8,15 +8,9 @@
 from nodestream_plugin_neo4j.query_executor import Neo4jQueryExecutor
 
 from .matchers import ran_query
-from .test_ingest_query_builder import (
-    DEFAULT_CHUNK_SIZE,
-    DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
-    DEFAULT_RETRIES_PER_CHUNK,
-)
+from .test_ingest_query_builder import DEFAULT_RETRIES_PER_CHUNK
 
 TEST_PARAMS = {
-    "execute_chunks_in_parallel": DEFAULT_EXECUTE_CHUNKS_IN_PARALLEL,
-    "chunk_size": DEFAULT_CHUNK_SIZE,
     "retries_per_chunk": DEFAULT_RETRIES_PER_CHUNK,
 }