Merge pull request #12 from nodestream-proj/ttl-fix

Batch TTL fix + TTL integration tests

Author: angelosantos4

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -254,12 +254,19 @@ 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,
+        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)
+        return ttl_match_query.feed_batched_query(
+            operation, config.batch_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,10 @@ 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,
+            retries_per_chunk=self.retries_per_chunk,
+        )
         await self.database_connection.execute(query)
 
     async def execute_hook(self, hook: IngestionHook):

tests/e2e/project/node-ttls.yaml

@@ -0,0 +1,8 @@
+- implementation: nodestream.pipeline.extractors:TimeToLiveConfigurationExtractor
+  arguments:
+    graph_object_type: NODE
+    configurations: 
+    - object_type: ObjectA
+      expiry_in_hours: 48
+    - object_type: ObjectB
+      expiry_in_hours: 24

tests/e2e/project/nodestream.yaml

@@ -5,6 +5,10 @@ scopes:
         name: airports
       - path: tests/e2e/project/fifa_2021_player_data.yaml
         name: fifa
+      - path: tests/e2e/project/node-ttls.yaml
+        name: node-ttls
+      - path: tests/e2e/project/relationship-ttls.yaml
+        name: relationship-ttls
 
 targets:
   my-neo4j-db:

tests/e2e/project/relationship-ttls.yaml

@@ -0,0 +1,8 @@
+- implementation: nodestream.pipeline.extractors:TimeToLiveConfigurationExtractor
+  arguments:
+    graph_object_type: RELATIONSHIP
+    configurations: 
+    - object_type: CONNECTED_TO
+      expiry_in_hours: 48
+    - object_type: ADJACENT_TO
+      expiry_in_hours: 24

tests/e2e/test_neo4j_pipelines.py

@@ -1,11 +1,15 @@
 from pathlib import Path
 
 import pytest
+from neo4j import Session
 from nodestream.pipeline import (
     PipelineInitializationArguments,
     PipelineProgressReporter,
 )
 from nodestream.project import Project, RunRequest
+from pandas import Timedelta, Timestamp
+
+from nodestream_plugin_neo4j.query import Query
 
 from .conftest import TESTED_NEO4J_VERSIONS
 
@@ -59,15 +63,212 @@ def validate_fifa_mo_club(session):
     assert result.single()["club"] == "Liverpool"
 
 
+def validate_consistency_in_node_counts(session):
+    result = session.run(
+        """
+        MATCH (n:ObjectA)
+        RETURN count(n) as node_count
+        """
+    )
+    assert result.single()["node_count"] == 30
+    result = session.run(
+        """
+        MATCH (n:ObjectB)
+        RETURN count(n) as node_count
+        """
+    )
+    assert result.single()["node_count"] == 30
+
+
+def validate_ttl_seperation_between_relationship_object_types(session):
+    result = session.run(
+        """
+        MATCH ()-[r:CONNECTED_TO]->()
+        RETURN count(r) as relationship_count
+        """
+    )
+    assert result.single()["relationship_count"] == 20
+    result = session.run(
+        """
+        MATCH ()-[r:ADJACENT_TO]->()
+        RETURN count(r) as relationship_count
+        """
+    )
+    assert result.single()["relationship_count"] == 10
+
+
+def validate_ttl_seperation_between_node_object_types(session):
+    result = session.run(
+        """
+        MATCH (n:ObjectA)
+        RETURN count(n) as node_count
+        """
+    )
+    assert result.single()["node_count"] == 20
+    result = session.run(
+        """
+        MATCH (n:ObjectB)
+        RETURN count(n) as node_count
+        """
+    )
+    assert result.single()["node_count"] == 10
+
+
+PIPELINE_TESTS = [
+    ("airports", [validate_airports, valiudate_airport_country]),
+    ("fifa", [validate_fifa_player_count, validate_fifa_mo_club]),
+]
+
+TTL_TESTS = [
+    (
+        "relationship-ttls",
+        [
+            validate_consistency_in_node_counts,
+            validate_ttl_seperation_between_relationship_object_types,
+        ],
+    ),
+    ("node-ttls", [validate_ttl_seperation_between_node_object_types]),
+]
+
+NODE_CREATION_QUERY = """
+CREATE (n:{node_label}) 
+SET n.last_ingested_at = $timestamp
+SET n.identifier = $node_id
+"""
+
+RELATIONSHIP_CREATION_QUERY = """
+MATCH (from_node:{from_node_label}) MATCH (to_node:{to_node_label}) 
+WHERE to_node.identifier=$to_node_identifier AND from_node.identifier=$from_node_identifier
+CREATE (from_node)-[rel:{relationship_label}]->(to_node) 
+SET rel.last_ingested_at = $timestamp
+"""
+
+REALLY_OLD_TIMESTAMP = Timestamp.utcnow() - Timedelta(hours=60)
+OLD_TIMESTAMP = Timestamp.utcnow() - Timedelta(hours=36)
+NEW_TIMESTAMP = Timestamp.utcnow() - Timedelta(hours=12)
+
+
+def create_node_query_from_params(node_label, node_id, timestamp):
+    return Query(
+        NODE_CREATION_QUERY.format(node_label=node_label),
+        {"timestamp": timestamp, "node_id": node_id},
+    )
+
+
+def create_relationship_query_from_params(
+    from_node_label,
+    from_node_identifier,
+    to_node_label,
+    to_node_identifier,
+    relationship_label,
+    timestamp,
+):
+    return Query(
+        RELATIONSHIP_CREATION_QUERY.format(
+            from_node_label=from_node_label,
+            to_node_label=to_node_label,
+            relationship_label=relationship_label,
+        ),
+        {
+            "timestamp": timestamp,
+            "from_node_identifier": from_node_identifier,
+            "to_node_identifier": to_node_identifier,
+        },
+    )
+
+
+"""
+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)
+        session.run(query.query_statement, query.parameters)
+
+    def create_relationship(
+        from_node_label,
+        from_node_identifier,
+        to_node_label,
+        to_node_identifier,
+        relationship_label,
+        timestamp,
+    ):
+        query = create_relationship_query_from_params(
+            from_node_label,
+            from_node_identifier,
+            to_node_label,
+            to_node_identifier,
+            relationship_label,
+            timestamp,
+        )
+        session.run(query.query_statement, query.parameters)
+
+    for i in range(0, 10):
+        create_node("ObjectA", str(i), REALLY_OLD_TIMESTAMP)
+        create_node("ObjectB", str(i), REALLY_OLD_TIMESTAMP)
+
+    for i in range(10, 20):
+        create_node("ObjectA", str(i), OLD_TIMESTAMP)
+        create_node("ObjectB", str(i), OLD_TIMESTAMP)
+
+    for i in range(20, 30):
+        create_node("ObjectA", str(i), NEW_TIMESTAMP)
+        create_node("ObjectB", str(i), NEW_TIMESTAMP)
+
+    for i in range(0, 10):
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "CONNECTED_TO", NEW_TIMESTAMP
+        )
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "ADJACENT_TO", NEW_TIMESTAMP
+        )
+
+    for i in range(10, 20):
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "CONNECTED_TO", OLD_TIMESTAMP
+        )
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "ADJACENT_TO", OLD_TIMESTAMP
+        )
+
+    for i in range(20, 30):
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "CONNECTED_TO", REALLY_OLD_TIMESTAMP
+        )
+        create_relationship(
+            "ObjectA", str(i), "ObjectB", str(i), "ADJACENT_TO", REALLY_OLD_TIMESTAMP
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.e2e
 @pytest.mark.parametrize("neo4j_version", TESTED_NEO4J_VERSIONS)
 @pytest.mark.parametrize(
     "pipeline_name,validations",
-    [
-        ("airports", [validate_airports, valiudate_airport_country]),
-        ("fifa", [validate_fifa_player_count, validate_fifa_mo_club]),
-    ],
+    PIPELINE_TESTS,
 )
 async def test_neo4j_pipeline(
     project, neo4j_container, pipeline_name, validations, neo4j_version
@@ -87,3 +288,24 @@ async def test_neo4j_pipeline(
 
         for validator in validations:
             validator(session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.e2e
+@pytest.mark.parametrize("neo4j_version", TESTED_NEO4J_VERSIONS)
+async def test_neo4j_ttls(project, neo4j_container, neo4j_version):
+    with neo4j_container(
+        neo4j_version
+    ) as neo4j_container, neo4j_container.get_driver() as driver, driver.session() as session:
+        create_test_objects(session)
+        target = project.get_target_by_name("my-neo4j-db")
+        for pipeline_name, validations in TTL_TESTS:
+            await project.run(
+                RunRequest(
+                    pipeline_name,
+                    PipelineInitializationArguments(extra_steps=[target.make_writer()]),
+                    PipelineProgressReporter(),
+                )
+            )
+            for validator in validations:
+                validator(session)