Fixes #257

Author: technige

neo4j/__init__.py

@@ -652,15 +652,16 @@ def _run_transaction(self, access_mode, unit_of_work, *args, **kwargs):
             raise TypeError("Unit of work is not callable")
 
         metadata = getattr(unit_of_work, "metadata", None)
-        print(metadata)    # TODO
+        timeout = getattr(unit_of_work, "timeout", None)
+
         retry_delay = retry_delay_generator(INITIAL_RETRY_DELAY,
                                             RETRY_DELAY_MULTIPLIER,
                                             RETRY_DELAY_JITTER_FACTOR)
         errors = []
         t0 = perf_counter()
         while True:
             try:
-                self._open_transaction(access_mode)
+                self._open_transaction(access_mode, metadata, timeout)
                 tx = self._transaction
                 try:
                     result = unit_of_work(tx, *args, **kwargs)
@@ -1405,7 +1406,7 @@ def __init__(self, transaction, *args, **kwargs):
         self.transaction = transaction
 
 
-def unit_of_work(**metadata):
+def unit_of_work(metadata=None, timeout=None):
     """ Decorator for transaction functions.
     """
 
@@ -1415,6 +1416,7 @@ def wrapped(*args, **kwargs):
             return f(*args, **kwargs)
 
         wrapped.metadata = metadata
+        wrapped.timeout = timeout
         return wrapped
 
     return wrapper

test/integration/test_session.py

@@ -579,14 +579,16 @@ def work(tx, x):
             self.assertEqual(value, 1)
 
     def test_read_with_arg_and_metadata(self):
+        if self.protocol_version() < 3:
+            raise SkipTest("Transaction metadata and timeout only supported in Bolt v3+")
 
         @unit_of_work(timeout=25, metadata={"foo": "bar"})
-        def work(tx, x):
-            return tx.run("RETURN $x", x=x).single().value()
+        def work(tx):
+            return tx.run("CALL dbms.getTXMetaData").single().value()
 
         with self.driver.session() as session:
-            value = session.read_transaction(work, x=1)
-            self.assertEqual(value, 1)
+            value = session.read_transaction(work)
+            self.assertEqual(value, {"foo": "bar"})
 
     def test_simple_write(self):
 
@@ -607,6 +609,8 @@ def work(tx, x):
             self.assertEqual(value, 1)
 
     def test_write_with_arg_and_metadata(self):
+        if self.protocol_version() < 3:
+            raise SkipTest("Transaction metadata and timeout only supported in Bolt v3+")
 
         @unit_of_work(timeout=25, metadata={"foo": "bar"})
         def work(tx, x):