Merge branch 'master' of https://github.com/wcmaier/redis-py into logging

3 files changed, 103 insertions, 8 deletions

diff --git a/redis/client.py b/redis/client.py
index 9ba3547..b80707e 100644
--- a/redis/client.py
+++ b/redis/client.py
@@ -1,5 +1,6 @@
 import datetime
 import errno
+import logging
 import socket
 import threading
 import time
@@ -8,6 +9,16 @@ from itertools import chain, imap
 from redis.exceptions import ConnectionError, ResponseError, InvalidResponse, WatchError
 from redis.exceptions import RedisError, AuthenticationError
 
+try:
+    NullHandler = logging.NullHandler
+except AttributeError:
+    class NullHandler(logging.Handler):
+        def emit(self, record): pass
+
+log = logging.getLogger("redis")
+# Add a no-op handler to avoid error messages if the importing module doesn't
+# configure logging.
+log.addHandler(NullHandler())
 
 class ConnectionPool(threading.local):
     "Manages a list of connections on the local thread"
@@ -47,6 +58,8 @@ class Connection(object):
         "Connects to the Redis server if not already connected"
         if self._sock:
             return
+        if log_enabled(log):
+            log.debug("connecting to %s:%d/%d", self.host, self.port, self.db)
         try:
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             sock.settimeout(self.socket_timeout)
@@ -70,6 +83,8 @@ class Connection(object):
         "Disconnects from the Redis server"
         if self._sock is None:
             return
+        if log_enabled(log):
+            log.debug("disconnecting from %s:%d/%d", self.host, self.port, self.db)
         try:
             self._sock.close()
         except socket.error:
@@ -148,6 +163,17 @@ def dict_merge(*dicts):
     [merged.update(d) for d in dicts]
     return merged
 
+def log_enabled(log, level=logging.DEBUG):
+    return log.isEnabledFor(log, level)
+
+def repr_command(args):
+    "Represents a command as a string."
+    command = [args[0]]
+    if len(args) > 1:
+        command.extend(repr(x) for x in args[1:])
+
+    return ' '.join(command)
+
 def parse_info(response):
     "Parse the result of Redis's INFO command into a Python dict"
     info = {}
@@ -317,6 +343,9 @@ class Redis(threading.local):
         if self.subscribed and not subscription_command:
             raise RedisError("Cannot issue commands other than SUBSCRIBE and "
                 "UNSUBSCRIBE while channels are open")
+        if log_enabled(log):
+            log.debug(repr_command(command))
+        command = self._encode_command(command)
         try:
             self.connection.send(command, self)
             if subscription_command:
@@ -328,14 +357,17 @@ class Redis(threading.local):
             if subscription_command:
                 return None
             return self.parse_response(command_name, **options)
+    
+    def _encode_command(self, args):
+        cmds = ['$%s\r\n%s\r\n' % (len(enc_value), enc_value)
+                for enc_value in imap(self.encode, args)]
+        return '*%s\r\n%s' % (len(cmds), ''.join(cmds))
 
     def execute_command(self, *args, **options):
         "Sends the command to the redis server and returns it's response"
-        cmds = ['$%s\r\n%s\r\n' % (len(enc_value), enc_value)
-                for enc_value in imap(self.encode, args)]
         return self._execute_command(
             args[0],
-            '*%s\r\n%s' % (len(cmds), ''.join(cmds)),
+            args,
             **options
             )
 
@@ -1407,11 +1439,16 @@ class Pipeline(Redis):
     def _execute_transaction(self, commands):
         # wrap the commands in MULTI ... EXEC statements to indicate an
         # atomic operation
-        all_cmds = ''.join([c for _1, c, _2 in chain(
-            (('', 'MULTI\r\n', ''),),
+        all_cmds = ''.join([self._encode_command(c) for _1, c, _2 in chain(
+            (('', ('MULTI',), ''),),
             commands,
-            (('', 'EXEC\r\n', ''),)
+            (('', ('EXEC',), ''),)
             )])
+        if log_enabled(log):
+            log.debug("MULTI")
+            for command in commands:
+                log.debug("TRANSACTION> "+ repr_command(command[1]))
+            log.debug("EXEC")
         self.connection.send(all_cmds, self)
         # parse off the response for MULTI and all commands prior to EXEC
         for i in range(len(commands)+1):
@@ -1436,7 +1473,10 @@ class Pipeline(Redis):
 
     def _execute_pipeline(self, commands):
         # build up all commands into a single request to increase network perf
-        all_cmds = ''.join([c for _1, c, _2 in commands])
+        all_cmds = ''.join([self._encode_command(c) for _1, c, _2 in commands])
+        if log_enabled(log):
+            for command in commands:
+                log.debug("PIPELINE> " + repr_command(command[1]))
         self.connection.send(all_cmds, self)
         data = []
         for command_name, _, options in commands:
diff --git a/tests/__init__.py b/tests/__init__.py
index 8931b07..63d6741 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,5 +1,5 @@
 import unittest
-from server_commands import ServerCommandsTestCase
+from server_commands import ServerCommandsTestCase, LoggingTestCase
 from connection_pool import ConnectionPoolTestCase
 from pipeline import PipelineTestCase
 from lock import LockTestCase
@@ -10,4 +10,5 @@ def all_tests():
     suite.addTest(unittest.makeSuite(ConnectionPoolTestCase))
     suite.addTest(unittest.makeSuite(PipelineTestCase))
     suite.addTest(unittest.makeSuite(LockTestCase))
+    suite.addTest(unittest.makeSuite(LoggingTestCase))
     return suite
diff --git a/tests/server_commands.py b/tests/server_commands.py
index 5cd63c5..f213321 100644
--- a/tests/server_commands.py
+++ b/tests/server_commands.py
@@ -3,6 +3,8 @@ import unittest
 import datetime
 import threading
 import time
+import logging
+import logging.handlers
 from distutils.version import StrictVersion
 
 class ServerCommandsTestCase(unittest.TestCase):
@@ -1258,3 +1260,55 @@ class ServerCommandsTestCase(unittest.TestCase):
         # check that it is possible to get list content by key name
         for key in mapping.keys():
             self.assertEqual(self.client.lrange(key, 0, -1), list(mapping[key]))
+
+class BufferingHandler(logging.handlers.BufferingHandler):
+
+    def __init__(self):
+        logging.handlers.BufferingHandler.__init__(self, None)
+
+    def shouldFlush(self, record):
+        return False
+
+class LoggingTestCase(unittest.TestCase):
+
+    def get_client(self):
+        return redis.Redis(host='localhost', port=6379, db=9)
+
+    def setUp(self):
+        self.client = self.get_client()
+        self.client.flushdb()
+
+        self.log = logging.getLogger("redis")
+        self.log.setLevel(logging.DEBUG)
+        self.handler = BufferingHandler()
+        self.log.addHandler(self.handler)
+        self.buffer = self.handler.buffer
+
+    def tearDown(self):
+        self.client.flushdb()
+        for c in self.client.connection_pool.get_all_connections():
+            c.disconnect()
+
+    def test_command_logging(self):
+        self.client.get("foo")
+
+        self.assertEqual(len(self.buffer), 1)
+        self.assertEqual(self.buffer[0].msg, "GET 'foo'")
+
+    def test_command_logging_pipeline(self):
+        pipe = self.client.pipeline(transaction=False)
+        pipe.get("foo")
+        pipe.execute()
+
+        self.assertEqual(len(self.buffer), 1)
+        self.assertEqual(self.buffer[0].msg, "PIPELINE> GET 'foo'")
+
+    def test_command_logging_transaction(self):
+        txn = self.client.pipeline(transaction=True)
+        txn.get("foo")
+        txn.execute()
+
+        self.assertEqual(len(self.buffer), 3)
+        messages = [x.msg for x in self.buffer]
+        self.assertEqual(messages,
+                ["MULTI", "TRANSACTION> GET 'foo'", "EXEC"])