Server address and version

Author: technige

neo4j/v1/bolt.py

@@ -29,7 +29,7 @@
 from __future__ import division
 
 from base64 import b64encode
-from collections import deque
+from collections import deque, namedtuple
 from io import BytesIO
 import logging
 from os import makedirs, open as os_open, write as os_write, close as os_close, O_CREAT, O_APPEND, O_WRONLY
@@ -81,12 +81,16 @@
 log_error = log.error
 
 
+Address = namedtuple("Address", ["host", "port"])
+ServerInfo = namedtuple("ServerInfo", ["address", "version"])
+
+
 class BufferingSocket(object):
 
     def __init__(self, connection):
         self.connection = connection
         self.socket = connection.socket
-        self.address = self.socket.getpeername()
+        self.address = Address(*self.socket.getpeername())
         self.buffer = bytearray()
 
     def fill(self):
@@ -132,7 +136,7 @@ class ChunkChannel(object):
 
     def __init__(self, sock):
         self.socket = sock
-        self.address = sock.getpeername()
+        self.address = Address(*sock.getpeername())
         self.raw = BytesIO()
         self.output_buffer = []
         self.output_size = 0
@@ -206,6 +210,22 @@ def on_ignored(self, metadata=None):
         pass
 
 
+class InitResponse(Response):
+
+    def on_success(self, metadata):
+        super(InitResponse, self).on_success(metadata)
+        connection = self.connection
+        address = Address(*connection.socket.getpeername())
+        version = metadata.get("server")
+        connection.server = ServerInfo(address, version)
+
+    def on_failure(self, metadata):
+        code = metadata.get("code")
+        error = (Unauthorized if code == "Neo.ClientError.Security.Unauthorized" else
+                 ServiceUnavailable)
+        raise error(metadata.get("message", "INIT failed"))
+
+
 class Connection(object):
     """ Server connection for Bolt protocol v1.
 
@@ -219,10 +239,12 @@ class Connection(object):
     #: The pool of which this connection is a member
     pool = None
 
+    #: Server version details
+    server = None
+
     def __init__(self, sock, **config):
         self.socket = sock
         self.buffering_socket = BufferingSocket(self)
-        self.address = sock.getpeername()
         self.channel = ChunkChannel(sock)
         self.packer = Packer(self.channel)
         self.unpacker = Unpacker()
@@ -246,15 +268,7 @@ 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_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_failure = on_failure
-
+        response = InitResponse(self)
         self.append(INIT, (self.user_agent, self.auth_dict), response=response)
         self.send()
         while not response.complete:
@@ -316,18 +330,18 @@ def send(self):
         """ Send all queued messages to the server.
         """
         if self.closed:
-            raise ServiceUnavailable("Failed to write to closed connection %r" % (self.address,))
+            raise ServiceUnavailable("Failed to write to closed connection %r" % (self.server.address,))
         if self.defunct:
-            raise ServiceUnavailable("Failed to write to defunct connection %r" % (self.address,))
+            raise ServiceUnavailable("Failed to write to defunct connection %r" % (self.server.address,))
         self.channel.send()
 
     def fetch(self):
         """ Receive exactly one message from the server.
         """
         if self.closed:
-            raise ServiceUnavailable("Failed to read from closed connection %r" % (self.address,))
+            raise ServiceUnavailable("Failed to read from closed connection %r" % (self.server.address,))
         if self.defunct:
-            raise ServiceUnavailable("Failed to read from defunct connection %r" % (self.address,))
+            raise ServiceUnavailable("Failed to read from defunct connection %r" % (self.server.address,))
         try:
             message_data = self.buffering_socket.read_message()
         except ProtocolError:

neo4j/v1/routing.py

@@ -22,7 +22,7 @@
 from threading import Lock
 from time import clock
 
-from .bolt import ConnectionPool
+from .bolt import Address, ConnectionPool
 from .compat.collections import MutableSet, OrderedDict
 from .exceptions import CypherError, ProtocolError, ServiceUnavailable
 
@@ -94,7 +94,7 @@ def parse_address(cls, address):
         """ Convert an address string to a tuple.
         """
         host, _, port = address.partition(":")
-        return host, int(port)
+        return Address(host, int(port))
 
     @classmethod
     def parse_routing_info(cls, records):

test/test_driver.py

@@ -159,7 +159,7 @@ def test_should_be_able_to_read(self):
                     result = session.run("RETURN $x", {"x": 1})
                     for record in result:
                         assert record["x"] == 1
-                    assert session.connection.address == ('127.0.0.1', 9004)
+                    assert session.connection.server.address == ('127.0.0.1', 9004)
 
     def test_should_be_able_to_write(self):
         with StubCluster({9001: "router.script", 9006: "create_a.script"}):
@@ -168,7 +168,7 @@ def test_should_be_able_to_write(self):
                 with driver.session(WRITE_ACCESS) as session:
                     result = session.run("CREATE (a $x)", {"x": {"name": "Alice"}})
                     assert not list(result)
-                    assert session.connection.address == ('127.0.0.1', 9006)
+                    assert session.connection.server.address == ('127.0.0.1', 9006)
 
     def test_should_be_able_to_write_as_default(self):
         with StubCluster({9001: "router.script", 9006: "create_a.script"}):
@@ -177,7 +177,7 @@ def test_should_be_able_to_write_as_default(self):
                 with driver.session() as session:
                     result = session.run("CREATE (a $x)", {"x": {"name": "Alice"}})
                     assert not list(result)
-                    assert session.connection.address == ('127.0.0.1', 9006)
+                    assert session.connection.server.address == ('127.0.0.1', 9006)
 
     def test_routing_disconnect_on_run(self):
         with StubCluster({9001: "router.script", 9004: "disconnect_on_run.script"}):

test/test_routing.py

@@ -575,7 +575,7 @@ def test_connected_to_reader(self):
             with RoutingConnectionPool(connector, address) as pool:
                 assert not pool.routing_table.is_fresh()
                 connection = pool.acquire_for_read()
-                assert connection.address in pool.routing_table.readers
+                assert connection.server.address in pool.routing_table.readers
 
     def test_should_retry_if_first_reader_fails(self):
         with StubCluster({9001: "router.script",
@@ -605,7 +605,7 @@ def test_connected_to_writer(self):
             with RoutingConnectionPool(connector, address) as pool:
                 assert not pool.routing_table.is_fresh()
                 connection = pool.acquire_for_write()
-                assert connection.address in pool.routing_table.writers
+                assert connection.server.address in pool.routing_table.writers
 
     def test_should_retry_if_first_writer_fails(self):
         with StubCluster({9001: "router_with_multiple_writers.script",