Fix bookmarks usage for retryable cypher

Author: FlorentinD

changelog.md

@@ -10,6 +10,7 @@
 ## Bug fixes
 
 - Fix reporting error based on http responses from the Aura-API with an invalid JSON body. Earlier the client would report JSONDecodeError instead of showing the actual issue.
+- Fixed a bug where retryable queries wouldnt work with bookmarks.
 
 ## Improvements
 

graphdatascience/query_runner/neo4j_query_runner.py

@@ -229,14 +229,20 @@ def run_retryable_cypher(
             routing = mode.neo4j_routing()
 
         try:
-            return self._driver.execute_query(
+            bookmark_manager = neo4j.GraphDatabase.bookmark_manager(self.bookmarks())
+
+            result = self._driver.execute_query(
                 query_=query,
                 parameters_=params,
                 database_=database,
                 result_transformer_=neo4j.Result.to_df,
-                bookmark_manager_=self.bookmarks(),
+                bookmark_manager_=bookmark_manager,
                 routing_=routing,
             )
+
+            self._last_bookmarks = neo4j.Bookmarks.from_raw_values(bookmark_manager.get_bookmarks())
+
+            return result
         except Exception as e:
             if custom_error:
                 Neo4jQueryRunner.handle_driver_exception(self._driver, e)

graphdatascience/tests/integration/test_database_ops.py

@@ -167,6 +167,25 @@ def test_run_read_cypher(gds: GraphDataScience) -> None:
     assert len(result) > 10
 
 
+def test_run_cypher_with_bookmarks(gds: GraphDataScience) -> None:
+    gds.run_cypher("RETURN true")
+    bookmarks_from_standard = gds.last_bookmarks()
+    gds.set_bookmarks(bookmarks_from_standard)
+
+    assert bookmarks_from_standard is not None
+
+    gds.run_cypher("RETURN true", retryable=True)
+    bookmarks_from_retryable = gds.last_bookmarks()
+
+    assert bookmarks_from_retryable is not None
+    assert bookmarks_from_retryable != bookmarks_from_standard
+
+    gds.set_bookmarks(bookmarks_from_retryable)
+
+    gds.run_cypher("RETURN true", retryable=True)
+    gds.run_cypher("RETURN true", retryable=False)
+
+
 def test_server_version(gds: GraphDataScience) -> None:
     cached_server_version = gds._server_version
     server_version_string = gds.version()