Improve Aura Example

tests/integration/examples/test_driver_introduction_example.py

@@ -25,8 +25,10 @@
 from io import StringIO
 
 # tag::driver-introduction-example-import[]
-from neo4j import GraphDatabase
 import logging
+import sys
+
+from neo4j import GraphDatabase
 from neo4j.exceptions import ServiceUnavailable
 # end::driver-introduction-example-import[]
 
@@ -45,27 +47,43 @@ def close(self):
         # Don't forget to close the driver connection when you are finished with it
         self.driver.close()
 
-    def create_friendship(self, person1_name, person2_name):
+    @staticmethod
+    def enable_log(level, output_stream):
+        handler = logging.StreamHandler(output_stream)
+        handler.setLevel(level)
+        logging.getLogger("neo4j").addHandler(handler)
+        logging.getLogger("neo4j").setLevel(level)
+
+    def create_friendship(self, person1_name, person2_name, knows_from):
         with self.driver.session() as session:
             # Write transactions allow the driver to handle retries and transient errors
             result = session.write_transaction(
-                self._create_and_return_friendship, person1_name, person2_name)
+                self._create_and_return_friendship, person1_name, person2_name, knows_from)
             for row in result:
-                print("Created friendship between: {p1}, {p2}".format(p1=row['p1'], p2=row['p2']))
+                print("Created friendship between: {p1}, {p2} from {knows_from}"
+                      .format(
+                          p1=row['p1'],
+                          p2=row['p2'],
+                          knows_from=row["knows_from"]))
 
     @staticmethod
-    def _create_and_return_friendship(tx, person1_name, person2_name):
+    def _create_and_return_friendship(tx, person1_name, person2_name, knows_from):
         # To learn more about the Cypher syntax, see https://neo4j.com/docs/cypher-manual/current/
         # The Reference Card is also a good resource for keywords https://neo4j.com/docs/cypher-refcard/current/
         query = (
             "CREATE (p1:Person { name: $person1_name }) "
             "CREATE (p2:Person { name: $person2_name }) "
-            "CREATE (p1)-[:KNOWS]->(p2) "
-            "RETURN p1, p2"
+            "CREATE (p1)-[k:KNOWS { from: $knows_from }]->(p2) "
+            "RETURN p1, p2, k"
         )
-        result = tx.run(query, person1_name=person1_name, person2_name=person2_name)
+        result = tx.run(query, person1_name=person1_name,
+                        person2_name=person2_name, knows_from=knows_from)
         try:
-            return [{"p1": row["p1"]["name"], "p2": row["p2"]["name"]}
+            return [{
+                        "p1": row["p1"]["name"],
+                        "p2": row["p2"]["name"],
+                        "knows_from": row["k"]["from"]
+                    }
                     for row in result]
         # Capture any errors along with the query and data for traceability
         except ServiceUnavailable as exception:
@@ -94,8 +112,9 @@ def _find_and_return_person(tx, person_name):
     bolt_url = "%%BOLT_URL_PLACEHOLDER%%"
     user = "<Username for database>"
     password = "<Password for database>"
+    App.enable_log(logging.INFO, sys.stdout)
     app = App(bolt_url, user, password)
-    app.create_friendship("Alice", "David")
+    app.create_friendship("Alice", "David", "School")
     app.find_person("Alice")
     app.close()
 # end::driver-introduction-example[]
@@ -105,8 +124,9 @@ def test_driver_introduction_example(uri, auth):
     try:
         s = StringIO()
         with redirect_stdout(s):
+            App.enable_log(logging.INFO, sys.stdout)
             app = App(uri, auth[0], auth[1])
-            app.create_friendship("Alice", "David")
+            app.create_friendship("Alice", "David", "School")
             app.find_person("Alice")
             app.close()