Proper reset handling

Author: technige

neo4j/__main__.py

@@ -26,53 +26,10 @@
 from json import loads as json_loads
 from sys import stdout, stderr
 
+from .util import Watcher
 from .v1.session import GraphDatabase, CypherError
 
 
-class ColourFormatter(logging.Formatter):
-    """ Colour formatter for pretty log output.
-    """
-
-    def format(self, record):
-        s = super(ColourFormatter, self).format(record)
-        if record.levelno == logging.CRITICAL:
-            return "\x1b[31;1m%s\x1b[0m" % s  # bright red
-        elif record.levelno == logging.ERROR:
-            return "\x1b[33;1m%s\x1b[0m" % s  # bright yellow
-        elif record.levelno == logging.WARNING:
-            return "\x1b[33m%s\x1b[0m" % s    # yellow
-        elif record.levelno == logging.INFO:
-            return "\x1b[36m%s\x1b[0m" % s    # cyan
-        elif record.levelno == logging.DEBUG:
-            return "\x1b[34m%s\x1b[0m" % s    # blue
-        else:
-            return s
-
-
-class Watcher(object):
-    """ Log watcher for debug output.
-    """
-
-    handlers = {}
-
-    def __init__(self, logger_name):
-        super(Watcher, self).__init__()
-        self.logger_name = logger_name
-        self.logger = logging.getLogger(self.logger_name)
-        self.formatter = ColourFormatter("%(asctime)s  %(message)s")
-
-    def watch(self, level=logging.INFO, out=stdout):
-        try:
-            self.logger.removeHandler(self.handlers[self.logger_name])
-        except KeyError:
-            pass
-        handler = logging.StreamHandler(out)
-        handler.setFormatter(self.formatter)
-        self.handlers[self.logger_name] = handler
-        self.logger.addHandler(handler)
-        self.logger.setLevel(level)
-
-
 def main():
     parser = ArgumentParser(description="Execute one or more Cypher statements using Bolt.")
     parser.add_argument("statement", nargs="+")

neo4j/util.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2016 "Neo Technology,"
+# Network Engine for Objects in Lund AB [http://neotechnology.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import unicode_literals
+
+import logging
+from argparse import ArgumentParser
+from json import loads as json_loads
+from sys import stdout, stderr
+
+from .v1.session import GraphDatabase, CypherError
+
+
+class ColourFormatter(logging.Formatter):
+    """ Colour formatter for pretty log output.
+    """
+
+    def format(self, record):
+        s = super(ColourFormatter, self).format(record)
+        if record.levelno == logging.CRITICAL:
+            return "\x1b[31;1m%s\x1b[0m" % s  # bright red
+        elif record.levelno == logging.ERROR:
+            return "\x1b[33;1m%s\x1b[0m" % s  # bright yellow
+        elif record.levelno == logging.WARNING:
+            return "\x1b[33m%s\x1b[0m" % s    # yellow
+        elif record.levelno == logging.INFO:
+            return "\x1b[36m%s\x1b[0m" % s    # cyan
+        elif record.levelno == logging.DEBUG:
+            return "\x1b[34m%s\x1b[0m" % s    # blue
+        else:
+            return s
+
+
+class Watcher(object):
+    """ Log watcher for debug output.
+    """
+
+    handlers = {}
+
+    def __init__(self, logger_name):
+        super(Watcher, self).__init__()
+        self.logger_name = logger_name
+        self.logger = logging.getLogger(self.logger_name)
+        self.formatter = ColourFormatter("%(asctime)s  %(message)s")
+
+    def watch(self, level=logging.INFO, out=stdout):
+        self.stop()
+        handler = logging.StreamHandler(out)
+        handler.setFormatter(self.formatter)
+        self.handlers[self.logger_name] = handler
+        self.logger.addHandler(handler)
+        self.logger.setLevel(level)
+
+    def stop(self):
+        try:
+            self.logger.removeHandler(self.handlers[self.logger_name])
+        except KeyError:
+            pass

neo4j/v1/connection.py

@@ -200,6 +200,10 @@ def on_ignored(self, metadata=None):
         pass
 
 
+class Completable(object):
+    complete = False
+
+
 class Connection(object):
     """ Server connection through which all protocol messages
     are sent and received. This class is designed for protocol
@@ -246,17 +250,22 @@ def append(self, signature, fields=(), response=None):
         self.channel.flush(end_of_message=True)
         self.responses.append(response)
 
-    def append_reset(self):
-        """ Add a RESET message to the outgoing queue.
+    def reset(self):
+        """ Add a RESET message to the outgoing queue, send
+        it and consume all remaining messages.
         """
+        response = Response(self)
 
         def on_failure(metadata):
             raise ProtocolError("Reset failed")
 
-        response = Response(self)
         response.on_failure = on_failure
 
         self.append(RESET, response=response)
+        self.send()
+        fetch_next = self.fetch_next
+        while not response.complete:
+            fetch_next()
 
     def send(self):
         """ Send all queued messages to the server.
@@ -280,18 +289,18 @@ def fetch_next(self):
         for signature, fields in unpack():
             if __debug__:
                 log_info("S: %s %s", message_names[signature], " ".join(map(repr, fields)))
+            if signature in SUMMARY:
+                response.complete = True
+                self.responses.popleft()
+            if signature == FAILURE:
+                self.reset()
             handler_name = "on_%s" % message_names[signature].lower()
             try:
                 handler = getattr(response, handler_name)
             except AttributeError:
                 pass
             else:
                 handler(*fields)
-            if signature in SUMMARY:
-                response.complete = True
-                self.responses.popleft()
-            if signature == FAILURE:
-                self.append_reset()
         raw.close()
 
     def close(self):

neo4j/v1/session.py

@@ -97,14 +97,18 @@ def session(self):
         """ Create a new session based on the graph database details
         specified within this driver:
 
+            >>> from neo4j.v1 import GraphDatabase
+            >>> driver = GraphDatabase.driver("bolt://localhost")
             >>> session = driver.session()
 
         """
         try:
             session = self.sessions.pop()
-            session.reset()
         except IndexError:
-            return Session(self)
+            session = Session(self)
+        else:
+            session.reset()
+        return session
 
 
 class Result(list):
@@ -123,7 +127,7 @@ def __init__(self, session, statement, parameters):
         self.statement = statement
         self.parameters = parameters
         self.keys = None
-        self.complete = False
+        self.more = True
         self.summary = None
         self.bench_test = None
 
@@ -142,7 +146,7 @@ def on_record(self, values):
     def on_footer(self, metadata):
         """ Called on receipt of the result footer.
         """
-        self.complete = True
+        self.more = False
         self.summary = ResultSummary(self.statement, self.parameters, **metadata)
         if self.bench_test:
             self.bench_test.end_recv = perf_counter()
@@ -157,7 +161,7 @@ def consume(self):
         callback functions.
         """
         fetch_next = self.session.connection.fetch_next
-        while not self.complete:
+        while self.more:
             fetch_next()
 
     def summarize(self):
@@ -340,6 +344,7 @@ def __init__(self, driver):
         self.connection = connect(driver.host, driver.port, **driver.config)
         self.transaction = None
         self.bench_tests = []
+        self.closed = False
 
     def __del__(self):
         self.connection.close()
@@ -353,7 +358,7 @@ def __exit__(self, exc_type, exc_value, traceback):
     def reset(self):
         """ Reset the connection so it can be reused from a clean state.
         """
-        self.connection.append_reset()
+        self.connection.reset()
 
     def run(self, statement, parameters=None):
         """ Run a parameterised Cypher statement.
@@ -407,9 +412,12 @@ def run(self, statement, parameters=None):
         return result
 
     def close(self):
-        """ Return this session to the driver pool it came from.
+        """ If still usable, return this session to the driver pool it came from.
         """
-        self.driver.sessions.appendleft(self)
+        self.reset()
+        if not self.connection.defunct:
+            self.driver.sessions.appendleft(self)
+        self.closed = True
 
     def begin_transaction(self):
         """ Create a new :class:`.Transaction` within this session.
@@ -487,6 +495,7 @@ def close(self):
         self.closed = True
         self.session.transaction = None
 
+
 class Record(object):
     """ Record is an ordered collection of fields.
 

test/test_session.py

@@ -23,6 +23,7 @@
 
 from neo4j.v1.session import GraphDatabase, CypherError, Record, record
 from neo4j.v1.typesystem import Node, Relationship, Path
+from test.util import watch
 
 
 class RunTestCase(TestCase):

test/util.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2016 "Neo Technology,"
+# Network Engine for Objects in Lund AB [http://neotechnology.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import functools
+
+from neo4j.util import Watcher
+
+
+def watch(f):
+    """ Decorator to enable log watching for the lifetime of a function.
+    Useful for debugging unit tests.
+
+    :param f: the function to decorate
+    :return: a decorated function
+    """
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        watcher = Watcher("neo4j")
+        watcher.watch()
+        f(*args, **kwargs)
+        watcher.stop()
+    return wrapper