Merge pull request #26 from nodestream-proj/beta/first-created-at-timestamp

Support Experimental `first_ingested_at` timestamp

Author: zprobst

nodestream_plugin_neo4j/database_connector.py

@@ -25,6 +25,7 @@ def from_file_data(
         chunk_size: int = 1000,
         execute_chunks_in_parallel: bool = True,
         retries_per_chunk: int = 3,
+        _experimental_set_first_ingested_at: bool = False,
         **connection_args
     ):
         database_connection = Neo4jDatabaseConnection.from_configuration(
@@ -37,6 +38,7 @@ def from_file_data(
             chunk_size=chunk_size,
             execute_chunks_in_parallel=execute_chunks_in_parallel,
             retries_per_chunk=retries_per_chunk,
+            set_first_ingested_at=_experimental_set_first_ingested_at,
         )
 
     def __init__(
@@ -47,16 +49,20 @@ def __init__(
         chunk_size: int = 1000,
         execute_chunks_in_parallel: bool = True,
         retries_per_chunk: int = 3,
+        set_first_ingested_at: bool = False,
     ) -> None:
         self.use_enterprise_features = use_enterprise_features
         self.use_apoc = use_apoc
         self.database_connection = database_connection
         self.chunk_size = chunk_size
         self.execute_chunks_in_parallel = execute_chunks_in_parallel
         self.retries_per_chunk = retries_per_chunk
+        self.set_first_ingested_at = set_first_ingested_at
 
     def make_query_executor(self) -> QueryExecutor:
-        query_builder = Neo4jIngestQueryBuilder(self.use_apoc)
+        query_builder = Neo4jIngestQueryBuilder(
+            self.use_apoc, self.set_first_ingested_at
+        )
         return Neo4jQueryExecutor(
             self.database_connection,
             query_builder,

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -94,7 +94,9 @@ def _merge_node(
 
 @cache
 def _make_relationship(
-    rel_identity: RelationshipIdentityShape, creation_rule: RelationshipCreationRule
+    rel_identity: RelationshipIdentityShape,
+    creation_rule: RelationshipCreationRule,
+    set_first_ingested_at: bool,
 ):
     keys = generate_properties_set_with_prefix(rel_identity.keys, RELATIONSHIP_REF_NAME)
     merge_rel_query = (
@@ -112,14 +114,20 @@ def _make_relationship(
     set_properties_query = f"SET {RELATIONSHIP_REF_NAME} += params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)}"
     if creation_rule == RelationshipCreationRule.CREATE:
         create_rel_query = str(merge_rel_query).replace("MERGE", "CREATE")
+        if set_first_ingested_at:
+            set_properties_query = f"{set_properties_query} SET {RELATIONSHIP_REF_NAME}.first_ingested_at = params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)}['last_ingested_at']"
         return f"{create_rel_query} {set_properties_query}"
 
+    if set_first_ingested_at:
+        merge_rel_query = f"{merge_rel_query} ON CREATE SET {RELATIONSHIP_REF_NAME}.first_ingested_at = params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)}['last_ingested_at']"
+
     return f"{merge_rel_query} {set_properties_query}"
 
 
 class Neo4jIngestQueryBuilder:
-    def __init__(self, apoc_iterate: bool):
+    def __init__(self, apoc_iterate: bool, set_first_ingested_at: bool):
         self.apoc_iterate = apoc_iterate
+        self.set_first_ingested_at = set_first_ingested_at
 
     @cache
     @correct_parameters
@@ -131,6 +139,12 @@ def generate_update_node_operation_query_statement(
 
         if operation.node_creation_rule == NodeCreationRule.EAGER:
             query = str(_merge_node(operation))
+
+            # We only need to add this if the user wants a first_ingested_at
+            # property  and we are in the merge case. Because MATCH will not
+            # create the node.
+            if self.set_first_ingested_at:
+                query = f"{query} ON CREATE SET {GENERIC_NODE_REF_NAME}.first_ingested_at = params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, GENERIC_NODE_REF_NAME)}['last_ingested_at']"
         else:
             query = str(_match_node(operation))
 
@@ -173,7 +187,9 @@ def generate_update_relationship_operation_query_statement(
         match_from_node_segment = _match_node(operation.from_node, FROM_NODE_REF_NAME)
         match_to_node_segment = _match_node(operation.to_node, TO_NODE_REF_NAME)
         merge_rel_segment = _make_relationship(
-            operation.relationship_identity, operation.relationship_creation_rule
+            operation.relationship_identity,
+            operation.relationship_creation_rule,
+            self.set_first_ingested_at,
         )
         return f"{match_from_node_segment} {match_to_node_segment} {merge_rel_segment}"
 

tests/e2e/project/nodestream.yaml

@@ -18,3 +18,9 @@ targets:
     uri: !env NEO4J_CONNECT_URI
     username: neo4j
     password: password
+  my-creation-ts-db:
+    database: neo4j
+    uri: !env NEO4J_CONNECT_URI
+    username: neo4j
+    password: password
+    _experimental_set_first_ingested_at: true

tests/e2e/test_neo4j_pipelines.py

@@ -19,6 +19,47 @@ def project():
     return Project.read_from_file(Path("tests/e2e/project/nodestream.yaml"))
 
 
+def validate_timestamps(session):
+    # Check that all nodes have a first_ingested_at property
+    nodes_with_first_ingested_at = session.run(
+        """
+        MATCH (n) where n.first_ingested_at IS NOT NULL
+        RETURN count(n) AS count
+        """
+    )
+
+    total_nodes = session.run(
+        """
+        MATCH (n)
+        RETURN count(n) AS count
+        """
+    )
+
+    assert (
+        nodes_with_first_ingested_at.single()["count"] == total_nodes.single()["count"]
+    )
+
+    # Check that all relationships have a last_ingested_at property
+    relationships_with_last_ingested_at = session.run(
+        """
+        MATCH ()-[r]->() where r.last_ingested_at IS NOT NULL
+        RETURN count(r) AS count
+        """
+    )
+
+    total_relationships = session.run(
+        """
+        MATCH ()-[r]->()
+        RETURN count(r) AS count
+        """
+    )
+
+    assert (
+        relationships_with_last_ingested_at.single()["count"]
+        == total_relationships.single()["count"]
+    )
+
+
 def validate_airports(session):
     result = session.run(
         """
@@ -125,8 +166,13 @@ def validate_ttl_seperation_between_node_object_types(session):
 
 
 PIPELINE_TESTS = [
-    ("airports", [validate_airports, valiudate_airport_country]),
-    ("fifa", [validate_fifa_player_count, validate_fifa_mo_club]),
+    ("airports", [validate_airports, valiudate_airport_country], "my-neo4j-db"),
+    ("fifa", [validate_fifa_player_count, validate_fifa_mo_club], "my-neo4j-db"),
+    (
+        "airports",
+        [validate_airports, valiudate_airport_country, validate_timestamps],
+        "my-creation-ts-db",
+    ),
 ]
 
 TTL_TESTS = [
@@ -280,16 +326,16 @@ def create_relationship(
 @pytest.mark.e2e
 @pytest.mark.parametrize("neo4j_version", TESTED_NEO4J_VERSIONS)
 @pytest.mark.parametrize(
-    "pipeline_name,validations",
+    "pipeline_name,validations,target_name",
     PIPELINE_TESTS,
 )
 async def test_neo4j_pipeline(
-    project, neo4j_container, pipeline_name, validations, neo4j_version
+    project, neo4j_container, pipeline_name, validations, neo4j_version, target_name
 ):
     with neo4j_container(
         neo4j_version
     ) as neo4j_container, neo4j_container.get_driver() as driver, driver.session() as session:
-        target = project.get_target_by_name("my-neo4j-db")
+        target = project.get_target_by_name(target_name)
 
         await project.run(
             RunRequest(

tests/unit/test_ingest_query_builder.py

@@ -27,7 +27,7 @@
 
 @pytest.fixture
 def query_builder():
-    return Neo4jIngestQueryBuilder(True)
+    return Neo4jIngestQueryBuilder(True, False)
 
 
 GREATEST_DAY = Timestamp(1998, 3, 25, 2, 0, 1)
@@ -177,6 +177,16 @@ def test_generates_expected_queries(mocked_utcnow, query_builder, ttl, expected_
         }
     ],
 )
+SIMPLE_NODE_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP = QueryBatch(
+    "MERGE (node: TestType {id : params.__node_id}) ON CREATE SET node.first_ingested_at = params.__node_properties['last_ingested_at'] SET node += params.__node_properties",
+    [
+        {
+            "__node_id": "foo",
+            "__node_properties": SIMPLE_NODE.properties,
+            "__node_additional_labels": (),
+        }
+    ],
+)
 
 SIMPLE_NODE_EXPECTED_QUERY_ON_MATCH = QueryBatch(
     "MATCH (node: TestType) WHERE node.id = params.__node_id SET node += params.__node_properties",
@@ -207,6 +217,17 @@ def test_generates_expected_queries(mocked_utcnow, query_builder, ttl, expected_
     ],
 )
 
+COMPLEX_NODE_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP = QueryBatch(
+    "MERGE (node: ComplexType {id : params.__node_id}) ON CREATE SET node.first_ingested_at = params.__node_properties['last_ingested_at'] WITH node, params CALL apoc.create.addLabels(node, params.__node_additional_labels) yield node as _ SET node += params.__node_properties",
+    [
+        {
+            "__node_id": "foo",
+            "__node_properties": COMPLEX_NODE.properties,
+            "__node_additional_labels": ("ExtraTypeOne", "ExtraTypeTwo"),
+        }
+    ],
+)
+
 COMPLEX_NODE_TWO = Node(
     "ComplexType",
     {"id_part1": "foo", "id_part2": "bar"},
@@ -225,6 +246,18 @@ def test_generates_expected_queries(mocked_utcnow, query_builder, ttl, expected_
     ],
 )
 
+COMPLEX_NODE_TWO_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP = QueryBatch(
+    "MERGE (node: ComplexType {id_part1 : params.__node_id_part1, id_part2 : params.__node_id_part2}) ON CREATE SET node.first_ingested_at = params.__node_properties['last_ingested_at'] WITH node, params CALL apoc.create.addLabels(node, params.__node_additional_labels) yield node as _ SET node += params.__node_properties",
+    [
+        {
+            "__node_id_part1": "foo",
+            "__node_id_part2": "bar",
+            "__node_properties": COMPLEX_NODE_TWO.properties,
+            "__node_additional_labels": ("ExtraTypeOne", "ExtraTypeTwo"),
+        }
+    ],
+)
+
 
 @pytest.mark.parametrize(
     "node,expected_query,node_creation_rule",
@@ -243,6 +276,36 @@ def test_node_update_generates_expected_queries(
     assert_that(query, equal_to(expected_query))
 
 
+@pytest.mark.parametrize(
+    "node,expected_query,node_creation_rule",
+    [
+        [
+            SIMPLE_NODE,
+            SIMPLE_NODE_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP,
+            NodeCreationRule.EAGER,
+        ],
+        [
+            COMPLEX_NODE,
+            COMPLEX_NODE_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP,
+            NodeCreationRule.EAGER,
+        ],
+        [
+            COMPLEX_NODE_TWO,
+            COMPLEX_NODE_TWO_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP,
+            NodeCreationRule.EAGER,
+        ],
+        [SIMPLE_NODE, SIMPLE_NODE_EXPECTED_QUERY_ON_MATCH, NodeCreationRule.MATCH_ONLY],
+    ],
+)
+def test_node_update_with_created_timestamp_generates_exepected_queries(
+    query_builder, node, expected_query, node_creation_rule
+):
+    query_builder.set_first_ingested_at = True
+    operation = OperationOnNodeIdentity(node.identity_shape, node_creation_rule)
+    query = query_builder.generate_batch_update_node_operation_batch(operation, [node])
+    assert_that(query, equal_to(expected_query))
+
+
 RELATIONSHIP_BETWEEN_TWO_NODES = RelationshipWithNodes(
     from_node=SIMPLE_NODE,
     to_node=COMPLEX_NODE,
@@ -262,6 +325,19 @@ def test_node_update_generates_expected_queries(
     ],
 )
 
+RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP = QueryBatch(
+    """MATCH (from_node: TestType) WHERE from_node.id = params.__from_node_id MATCH (to_node: ComplexType) WHERE to_node.id = params.__to_node_id
+    MERGE (from_node)-[rel: RELATED_TO]->(to_node) ON CREATE SET rel.first_ingested_at = params.__rel_properties['last_ingested_at'] SET rel += params.__rel_properties
+    """,
+    [
+        {
+            "__from_node_id": "foo",
+            "__to_node_id": "foo",
+            "__rel_properties": RELATIONSHIP_BETWEEN_TWO_NODES.relationship.properties,
+        }
+    ],
+)
+
 RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY = RelationshipWithNodes(
     from_node=SIMPLE_NODE,
     to_node=COMPLEX_NODE_TWO,
@@ -281,6 +357,21 @@ def test_node_update_generates_expected_queries(
     ],
 )
 
+
+RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_MULTI_KEY_AND_TIMESTAMP = QueryBatch(
+    """MATCH (from_node: TestType) WHERE from_node.id = params.__from_node_id MATCH (to_node: ComplexType) WHERE to_node.id_part1 = params.__to_node_id_part1 AND to_node.id_part2 = params.__to_node_id_part2
+    MERGE (from_node)-[rel: RELATED_TO]->(to_node) ON CREATE SET rel.first_ingested_at = params.__rel_properties['last_ingested_at'] SET rel += params.__rel_properties""",
+    [
+        {
+            "__from_node_id": "foo",
+            "__to_node_id_part1": "foo",
+            "__to_node_id_part2": "bar",
+            "__rel_properties": RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY.relationship.properties,
+        }
+    ],
+)
+
+
 RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY_AND_CREATE = RelationshipWithNodes(
     from_node=SIMPLE_NODE,
     to_node=COMPLEX_NODE_TWO,
@@ -301,6 +392,19 @@ def test_node_update_generates_expected_queries(
     ],
 )
 
+RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_MULTI_KEY_AND_CREATE_AND_TIMESTAMP = QueryBatch(
+    """MATCH (from_node: TestType) WHERE from_node.id = params.__from_node_id MATCH (to_node: ComplexType) WHERE to_node.id_part1 = params.__to_node_id_part1 AND to_node.id_part2 = params.__to_node_id_part2
+    CREATE (from_node)-[rel: RELATED_TO]->(to_node) SET rel += params.__rel_properties SET rel.first_ingested_at = params.__rel_properties['last_ingested_at']""",
+    [
+        {
+            "__from_node_id": "foo",
+            "__to_node_id_part1": "foo",
+            "__to_node_id_part2": "bar",
+            "__rel_properties": RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY_AND_CREATE.relationship.properties,
+        }
+    ],
+)
+
 
 @pytest.mark.parametrize(
     "rel,expected_query",
@@ -339,3 +443,46 @@ def test_relationship_update_generates_expected_queries(
     assert_that(
         query.batched_parameter_sets, equal_to(expected_query.batched_parameter_sets)
     )
+
+
+@pytest.mark.parametrize(
+    "rel,expected_query",
+    [
+        [
+            RELATIONSHIP_BETWEEN_TWO_NODES,
+            RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_CREATED_TIMESTAMP,
+        ],
+        [
+            RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY,
+            RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_MULTI_KEY_AND_TIMESTAMP,
+        ],
+        [
+            RELATIONSHIP_BETWEEN_TWO_NODES_WITH_MULTI_KEY_AND_CREATE,
+            RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_MULTI_KEY_AND_CREATE_AND_TIMESTAMP,
+        ],
+    ],
+)
+def test_relationship_update_with_created_timestamp_generates_expected_queries(
+    query_builder, rel, expected_query
+):
+    query_builder.set_first_ingested_at = True
+    to_op = OperationOnNodeIdentity(rel.to_node.identity_shape, NodeCreationRule.EAGER)
+    from_op = OperationOnNodeIdentity(
+        rel.from_node.identity_shape, NodeCreationRule.MATCH_ONLY
+    )
+    operation = OperationOnRelationshipIdentity(
+        from_op,
+        to_op,
+        rel.relationship.identity_shape,
+        relationship_creation_rule=rel.relationship_creation_rule,
+    )
+    query = query_builder.generate_batch_update_relationship_query_batch(
+        operation, [rel]
+    )
+    assert_that(
+        query.query_statement,
+        equal_to_ignoring_whitespace(expected_query.query_statement),
+    )
+    assert_that(
+        query.batched_parameter_sets, equal_to(expected_query.batched_parameter_sets)
+    )