Added db to RUN message

Author: technige

neo4j/io/__init__.py

@@ -163,8 +163,20 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
-    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
-        raise NotImplementedError
+    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None,
+            timeout=None, db=None, **handlers):
+        """ Appends a RUN message to the output stream.
+
+        :param statement: Cypher query string
+        :param parameters: dictionary of Cypher parameters
+        :param mode: access mode for routing - "READ" or "WRITE" (default)
+        :param bookmarks: iterable of bookmark values after which this transaction should begin
+        :param metadata: custom metadata dictionary to attach to the transaction
+        :param timeout: timeout for transaction execution (seconds)
+        :param db: name of the database against which to begin the transaction
+        :param handlers: handler functions passed into the returned Response object
+        :return: :class:`neo4j.io._bolt3.Response` object
+        """
 
     def discard_all(self, **handlers):
         raise NotImplementedError

neo4j/io/_bolt3.py

@@ -142,7 +142,10 @@ def hello(self):
         self.send_all()
         self.fetch_all()
 
-    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
+    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None,
+            timeout=None, db=None, **handlers):
+        if db is not None:
+            raise ValueError("Database selection is not supported in Bolt 3")
         if not parameters:
             parameters = {}
         extra = {}

neo4j/io/_bolt4x0.py

@@ -142,12 +142,15 @@ def hello(self):
         self.send_all()
         self.fetch_all()
 
-    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
+    def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=None,
+            timeout=None, db=None, **handlers):
         if not parameters:
             parameters = {}
         extra = {}
         if mode:
             extra["mode"] = mode
+        if db:
+            extra["db"] = db
         if bookmarks:
             try:
                 extra["bookmarks"] = list(bookmarks)

neo4j/work/simple.py

@@ -200,9 +200,12 @@ def done(summary_metadata):
                 raise ValueError("Metadata can only be attached at transaction level")
             if statement_timeout:
                 raise ValueError("Timeouts only apply at transaction level")
+            # TODO: fail if explicit database name has been set
         else:
             run_metadata["bookmarks"] = self._bookmarks_in
 
+        # TODO: capture ValueError and surface as SessionError/TransactionError if
+        # TODO: explicit database selection has been made
         cx.run(statement_text, parameters, **run_metadata)
         cx.pull_all(
             on_records=lambda records: result._records.extend(

tests/unit/io/test_class_bolt3.py

@@ -42,8 +42,15 @@ def test_conn_not_timed_out(fake_socket):
     assert connection.timedout() is False
 
 
-def test_db_extra_not_supported(fake_socket):
+def test_db_extra_not_supported_in_begin(fake_socket):
     address = ("127.0.0.1", 7687)
     connection = Bolt3(address, fake_socket(address))
     with pytest.raises(ValueError):
         connection.begin(db="something")
+
+
+def test_db_extra_not_supported_in_run(fake_socket):
+    address = ("127.0.0.1", 7687)
+    connection = Bolt3(address, fake_socket(address))
+    with pytest.raises(ValueError):
+        connection.run("", db="something")

tests/unit/io/test_class_bolt4x0.py

@@ -40,7 +40,7 @@ def test_conn_not_timed_out(fake_socket):
     assert connection.timedout() is False
 
 
-def test_db_extra(fake_socket):
+def test_db_extra_in_begin(fake_socket):
     address = ("127.0.0.1", 7687)
     socket = fake_socket(address)
     connection = Bolt4x0(address, socket)
@@ -50,3 +50,17 @@ def test_db_extra(fake_socket):
     assert tag == b"\x11"
     assert len(fields) == 1
     assert fields[0] == {"db": "something"}
+
+
+def test_db_extra_in_run(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.run("", {}, db="something")
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x10"
+    assert len(fields) == 3
+    assert fields[0] == ""
+    assert fields[1] == {}
+    assert fields[2] == {"db": "something"}