Send BEGIN eagerly on transaction begin

Author: robsdedude

neo4j/io/__init__.py

@@ -29,6 +29,7 @@
 __all__ = [
     "Bolt",
     "BoltPool",
+    "ConnectionErrorHandler",
     "Neo4jPool",
     "check_supported_server_product",
 ]
@@ -92,6 +93,7 @@
 )
 from neo4j.io._common import (
     CommitResponse,
+    ConnectionErrorHandler,
     Inbox,
     InitResponse,
     Outbox,

neo4j/io/_common.py

@@ -139,6 +139,43 @@ def view(self):
             return memoryview(self._data[:end])
 
 
+class ConnectionErrorHandler:
+    """
+    Wrapper class for handling connection errors.
+
+    The class will wrap each method to invoke a callback if the method raises
+    Neo4jError, SessionExpired, or ServiceUnavailable.
+    The error will be re-raised after the callback.
+    """
+
+    def __init__(self, connection, on_error):
+        """
+        :param connection the connection object to warp
+        :type connection Bolt
+        :param on_error the function to be called when a method of
+            connection raises of of the caught errors.
+        :type on_error callable
+        """
+        self.__connection = connection
+        self.__on_error = on_error
+
+    def __getattr__(self, item):
+        connection_attr = getattr(self.__connection, item)
+        if not callable(connection_attr):
+            return connection_attr
+
+        def outer(func):
+            def inner(*args, **kwargs):
+                try:
+                    func(*args, **kwargs)
+                except (Neo4jError, ServiceUnavailable, SessionExpired) as exc:
+                    self.__on_error(exc)
+                    raise
+            return inner
+
+        return outer(connection_attr)
+
+
 class Response:
     """ Subscriber object for a full response (zero or
     more detail messages followed by one summary message).

neo4j/work/result.py

@@ -23,51 +23,10 @@
 from warnings import warn
 
 from neo4j.data import DataDehydrator
-from neo4j.exceptions import (
-    Neo4jError,
-    ServiceUnavailable,
-    SessionExpired,
-)
+from neo4j.io import ConnectionErrorHandler
 from neo4j.work.summary import ResultSummary
 
 
-class _ConnectionErrorHandler:
-    """
-    Wrapper class for handling connection errors.
-
-    The class will wrap each method to invoke a callback if the method raises
-    SessionExpired or ServiceUnavailable.
-    The error will be re-raised after the callback.
-    """
-
-    def __init__(self, connection, on_error):
-        """
-        :param connection the connection object to warp
-        :type connection Bolt
-        :param on_error the function to be called when a method of
-            connection raises of of the caught errors.
-        :type on_error callable
-        """
-        self.connection = connection
-        self.on_error = on_error
-
-    def __getattr__(self, item):
-        connection_attr = getattr(self.connection, item)
-        if not callable(connection_attr):
-            return connection_attr
-
-        def outer(func):
-            def inner(*args, **kwargs):
-                try:
-                    func(*args, **kwargs)
-                except (Neo4jError, ServiceUnavailable, SessionExpired) as exc:
-                    self.on_error(exc)
-                    raise
-            return inner
-
-        return outer(connection_attr)
-
-
 class Result:
     """A handler for the result of Cypher query execution. Instances
     of this class are typically constructed and returned by
@@ -76,7 +35,7 @@ class Result:
 
     def __init__(self, connection, hydrant, fetch_size, on_closed,
                  on_error):
-        self._connection = _ConnectionErrorHandler(connection, on_error)
+        self._connection = ConnectionErrorHandler(connection, on_error)
         self._hydrant = hydrant
         self._on_closed = on_closed
         self._metadata = None
@@ -98,7 +57,7 @@ def __init__(self, connection, hydrant, fetch_size, on_closed,
 
     @property
     def _qid(self):
-        if self._raw_qid == self._connection.connection.most_recent_qid:
+        if self._raw_qid == self._connection.most_recent_qid:
             return -1
         else:
             return self._raw_qid
@@ -127,7 +86,7 @@ def on_attached(metadata):
             # For auto-commit there is no qid and Bolt 3 does not support qid
             self._raw_qid = metadata.get("qid", -1)
             if self._raw_qid != -1:
-                self._connection.connection.most_recent_qid = self._raw_qid
+                self._connection.most_recent_qid = self._raw_qid
             self._keys = metadata.get("fields")
             self._attached = True
 

neo4j/work/transaction.py

@@ -22,10 +22,9 @@
 from neo4j.work.result import Result
 from neo4j.data import DataHydrator
 from neo4j.exceptions import (
-    ServiceUnavailable,
-    SessionExpired,
     TransactionError,
 )
+from neo4j.io import ConnectionErrorHandler
 
 
 class Transaction:
@@ -41,6 +40,9 @@ class Transaction:
 
     def __init__(self, connection, fetch_size, on_closed, on_error):
         self._connection = connection
+        self._error_handling_connection = ConnectionErrorHandler(
+            connection, self._error_handler
+        )
         self._bookmark = None
         self._results = []
         self._closed = False
@@ -63,14 +65,15 @@ def __exit__(self, exception_type, exception_value, traceback):
     def _begin(self, database, bookmarks, access_mode, metadata, timeout):
         self._connection.begin(bookmarks=bookmarks, metadata=metadata,
                                timeout=timeout, mode=access_mode, db=database)
+        self._error_handling_connection.send_all()
+        self._error_handling_connection.fetch_all()
 
     def _result_on_closed_handler(self):
         pass
 
-    def _result_on_error_handler(self, exc):
+    def _error_handler(self, exc):
         self._last_error = exc
-        if isinstance(exc, (ServiceUnavailable, SessionExpired)):
-            self._closed = True
+        self._closed = True
         self._on_error(exc)
 
     def _consume_results(self):
@@ -126,7 +129,7 @@ def run(self, query, parameters=None, **kwparameters):
         result = Result(
             self._connection, DataHydrator(), self._fetch_size,
             self._result_on_closed_handler,
-            self._result_on_error_handler
+            self._error_handler
         )
         self._results.append(result)