Merge pull request #14 from nodestream-proj/neo4j_reader_fix

Neo4j Reader Fix + Integration Test

Author: zprobst

nodestream_plugin_neo4j/neo4j_database.py

@@ -27,7 +27,7 @@ def __init__(self, driver: AsyncDriver, database_name: str) -> None:
         self.logger = getLogger(self.__class__.__name__)
 
     async def execute(
-        self, query: Query, log_result: bool = False, routing=RoutingControl.WRITE
+        self, query: Query, log_result: bool = False, routing_=RoutingControl.WRITE
     ) -> Iterable[Record]:
         self.logger.info(
             "Executing Cypher Query to Neo4j",
@@ -41,7 +41,7 @@ async def execute(
             query.query_statement,
             query.parameters,
             database_=self.database_name,
-            routing_=routing,
+            routing_=routing_,
         )
         if log_result:
             for record in result.records:

tests/e2e/project/extractor_integration.yaml

@@ -0,0 +1,18 @@
+- implementation: nodestream_plugin_neo4j.extractor:Neo4jExtractor
+  arguments:
+    database: neo4j
+    uri: !env NEO4J_CONNECT_URI
+    username: neo4j
+    password: password
+    query: |
+      MATCH (a:ObjectA)
+      RETURN a.identifier as identifier
+      SKIP $offset LIMIT $limit
+
+- implementation: nodestream.interpreting:Interpreter
+  arguments:
+    interpretations:
+      - type: source_node
+        node_type: NewObject
+        key:
+          identifier: !jmespath identifier

tests/e2e/project/nodestream.yaml

@@ -9,6 +9,8 @@ scopes:
         name: node-ttls
       - path: tests/e2e/project/relationship-ttls.yaml
         name: relationship-ttls
+      - path: tests/e2e/project/extractor_integration.yaml
+        name: extractor_integration
 
 targets:
   my-neo4j-db:

tests/e2e/test_neo4j_pipelines.py

@@ -30,6 +30,16 @@ def validate_airports(session):
     assert result.single()["count"] == 1000
 
 
+def validate_enrichment(session):
+    result = session.run(
+        """
+        MATCH (a:NewObject)
+        RETURN count(a) AS count
+        """
+    )
+    assert result.single()["count"] == 30
+
+
 def valiudate_airport_country(session):
     result = session.run(
         """
@@ -130,6 +140,9 @@ def validate_ttl_seperation_between_node_object_types(session):
     ("node-ttls", [validate_ttl_seperation_between_node_object_types]),
 ]
 
+EXTRACTOR_TESTS = [("extractor_integration", [validate_enrichment])]
+
+
 NODE_CREATION_QUERY = """
 CREATE (n:{node_label}) 
 SET n.last_ingested_at = $timestamp
@@ -309,3 +322,24 @@ async def test_neo4j_ttls(project, neo4j_container, neo4j_version):
             )
             for validator in validations:
                 validator(session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.e2e
+@pytest.mark.parametrize("neo4j_version", TESTED_NEO4J_VERSIONS)
+async def test_neo4j_extractor(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 EXTRACTOR_TESTS:
+            await project.run(
+                RunRequest(
+                    pipeline_name,
+                    PipelineInitializationArguments(extra_steps=[target.make_writer()]),
+                    PipelineProgressReporter(),
+                )
+            )
+            for validator in validations:
+                validator(session)