Bookmarking

Author: technige

README.rst

@@ -33,7 +33,6 @@ Example Usage
         with session.begin_transaction() as write_tx:
             write_tx.run("CREATE (a:Person {name:{name},age:{age}})", name="Alice", age=33)
             write_tx.run("CREATE (a:Person {name:{name},age:{age}})", name="Bob", age=44)
-            write_tx.success = True
 
         with session.begin_transaction() as read_tx:
             result = read_tx.run("MATCH (a:Person) RETURN a.name AS name, a.age AS age")

neo4j/v1/bolt.py

@@ -211,6 +211,14 @@ class Connection(object):
     .. note:: logs at INFO level
     """
 
+    in_use = False
+
+    closed = False
+
+    defunct = False
+
+    server_version = None
+
     def __init__(self, sock, **config):
         self.socket = sock
         self.buffering_socket = BufferingSocket(sock)
@@ -219,9 +227,6 @@ def __init__(self, sock, **config):
         self.packer = Packer(self.channel)
         self.unpacker = Unpacker()
         self.responses = deque()
-        self.in_use = False
-        self.closed = False
-        self.defunct = False
 
         # Determine the user agent and ensure it is a Unicode value
         user_agent = config.get("user_agent", DEFAULT_USER_AGENT)
@@ -238,19 +243,21 @@ def __init__(self, sock, **config):
         # Pick up the server certificate, if any
         self.der_encoded_server_certificate = config.get("der_encoded_server_certificate")
 
+        def on_success(metadata):
+            self.server_version = metadata.get("server")
+
         def on_failure(metadata):
             code = metadata.get("code")
             error = (Unauthorized if code == "Neo.ClientError.Security.Unauthorized" else
                      ServiceUnavailable)
             raise error(metadata.get("message", "INIT failed"))
 
         response = Response(self)
+        response.on_success = on_success
         response.on_failure = on_failure
 
         self.append(INIT, (self.user_agent, self.auth_dict), response=response)
-        self.send()
-        while not response.complete:
-            self.fetch()
+        self.sync()
 
     def __del__(self):
         self.close()
@@ -359,7 +366,10 @@ def fetch(self):
         else:
             raise ProtocolError("Unexpected response message with signature %02X" % signature)
 
-    def fetch_all(self):
+    def sync(self):
+        """ Send and fetch all outstanding messages.
+        """
+        self.send()
         while self.responses:
             response = self.responses[0]
             while not response.complete:

neo4j/v1/session.py

@@ -311,6 +311,8 @@ class Session(object):
 
     transaction = None
 
+    last_bookmark = None
+
     def __init__(self, driver, connection, access_mode=None):
         self.driver = driver
         self.connection = connection
@@ -336,6 +338,8 @@ def run(self, statement, parameters=None, **kwparameters):
         :return: Cypher result
         :rtype: :class:`.StatementResult`
         """
+        self.last_bookmark = None
+
         statement = _norm_statement(statement)
         parameters = _norm_parameters(parameters, **kwparameters)
 
@@ -372,11 +376,11 @@ def close(self):
             self.transaction.close()
         if self.connection:
             if not self.connection.closed:
-                self.connection.fetch_all()
+                self.connection.sync()
             self.connection.in_use = False
             self.connection = None
 
-    def begin_transaction(self):
+    def begin_transaction(self, bookmark=None):
         """ Create a new :class:`.Transaction` within this session.
 
         :return: new :class:`.Transaction` instance.
@@ -387,15 +391,23 @@ def begin_transaction(self):
         def clear_transaction():
             self.transaction = None
 
-        self.run("BEGIN")
+        parameters = {}
+        if bookmark is not None:
+            parameters["bookmark"] = bookmark
+
+        self.run("BEGIN", parameters)
         self.transaction = Transaction(self, on_close=clear_transaction)
         return self.transaction
 
     def commit_transaction(self):
-        self.run("COMMIT")
+        result = self.run("COMMIT")
+        self.connection.sync()
+        summary = result.summary()
+        self.last_bookmark = summary.metadata.get("bookmark")
 
     def rollback_transaction(self):
         self.run("ROLLBACK")
+        self.connection.sync()
 
 
 class Transaction(object):
@@ -413,7 +425,7 @@ class Transaction(object):
     #: and rolled back otherwise. This attribute can be set in user code
     #: multiple times before a transaction completes with only the final
     #: value taking effect.
-    success = False
+    success = None
 
     #: Indicator to show whether the transaction has been closed, either
     #: with commit or rollback.
@@ -427,8 +439,8 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        if exc_value:
-            self.success = False
+        if self.success is None:
+            self.success = not bool(exc_type)
         self.close()
 
     def run(self, statement, parameters=None, **kwparameters):

test/test_session.py

@@ -19,7 +19,10 @@
 # limitations under the License.
 
 from mock import patch
+from unittest import skipUnless
+from uuid import uuid4
 
+from neo4j.v1 import READ_ACCESS, WRITE_ACCESS
 from neo4j.v1.exceptions import CypherError, ResultError
 from neo4j.v1.session import GraphDatabase, basic_auth, Record
 from neo4j.v1.types import Node, Relationship, Path
@@ -31,6 +34,21 @@
 AUTH_TOKEN = basic_auth("neotest", "neotest")
 
 
+def get_server_version():
+    driver = GraphDatabase.driver(BOLT_URI, auth=AUTH_TOKEN, encrypted=False)
+    with driver.session() as session:
+        full_version = session.connection.server_version
+        if full_version is None:
+            return "Neo4j", (3, 0), ()
+        product, _, tagged_version = full_version.partition("/")
+        tags = tagged_version.split("-")
+        version = map(int, tags[0].split("."))
+        return product, tuple(version), tuple(tags[1:])
+
+
+SERVER_PRODUCT, SERVER_VERSION, SERVER_TAGS = get_server_version()
+
+
 class AutoCommitTransactionTestCase(ServerTestCase):
 
     def setUp(self):
@@ -415,12 +433,52 @@ def test_can_rollback_transaction_using_with_block(self):
                 tx.run("MATCH (a) WHERE id(a) = {n} "
                        "SET a.foo = {foo}", {"n": node_id, "foo": "bar"})
 
+                tx.success = False
+
             # Check the property value
             result = session.run("MATCH (a) WHERE id(a) = {n} "
                                  "RETURN a.foo", {"n": node_id})
             assert len(list(result)) == 0
 
 
+class BookmarkingTestCase(ServerTestCase):
+
+    def setUp(self):
+        self.driver = GraphDatabase.driver(BOLT_URI, auth=AUTH_TOKEN, encrypted=False)
+
+    def tearDown(self):
+        self.driver.close()
+
+    @skipUnless(SERVER_VERSION >= (3, 1), "Bookmarking is not supported by this version of Neo4j")
+    def test_can_obtain_bookmark_after_commit(self):
+        with self.driver.session() as session:
+            with session.begin_transaction() as tx:
+                tx.run("RETURN 1")
+            assert session.last_bookmark is not None
+
+    @skipUnless(SERVER_VERSION >= (3, 1), "Bookmarking is not supported by this version of Neo4j")
+    def test_can_pass_bookmark_into_next_transaction(self):
+        unique_id = uuid4().hex
+
+        with self.driver.session(WRITE_ACCESS) as session:
+            with session.begin_transaction() as tx:
+                tx.run("CREATE (a:Thing {uuid:$uuid})", uuid=unique_id)
+            bookmark = session.last_bookmark
+
+        assert bookmark is not None
+
+        with self.driver.session(READ_ACCESS) as session:
+            with session.begin_transaction(bookmark) as tx:
+                result = tx.run("MATCH (a:Thing {uuid:$uuid}) RETURN a", uuid=unique_id)
+                record_list = list(result)
+                assert len(record_list) == 1
+                record = record_list[0]
+                assert len(record) == 1
+                thing = record[0]
+                assert isinstance(thing, Node)
+                assert thing["uuid"] == unique_id
+
+
 class ResultConsumptionTestCase(ServerTestCase):
 
     def setUp(self):