Merge branch 'master' into 2.0

Conflicts: cassandra/connection.py requirements.txt setup.py tests/unit/test_connection.py tox.ini
2014-05-07 17:06:13 -05:00
parent 4a702209fd 9206341d9d
commit 80715767a6
9 changed files with 194 additions and 27 deletions
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -43,6 +43,11 @@ now:
 =====
 In Progress

+Features
+--------
+* Allow a specific compression type to be requested for communications with
+  Cassandra and prefer lz4 if available
+
 Bug Fixes
 ---------
 * Update token metadata (for TokenAware calculations) when a node is removed
@@ -60,6 +65,7 @@ Other
 -----
 * Don't log at ERROR when a connection is closed during the startup
  communications
+* Mke scales, blist optional dependencies

 1.1.1
 =====
--- a/cassandra/cluster.py
+++ b/cassandra/cluster.py
@@ -185,8 +185,15 @@ class Cluster(object):

    compression = True
    """
-    Whether or not compression should be enabled when possible. Defaults to
-    :const:`True` and attempts to use snappy compression.
+    Controls compression for communications between the driver and Cassandra.
+    If left as the default of :const:`True`, either lz4 or snappy compression
+    may be used, depending on what is supported by both the driver
+    and Cassandra.  If both are fully supported, lz4 will be preferred.
+
+    You may also set this to 'snappy' or 'lz4' to request that specific
+    compression type.
+
+    Setting this to :const:`False` disables compression.
    """

    auth_provider = None
--- a/cassandra/connection.py
+++ b/cassandra/connection.py
@@ -25,6 +25,7 @@ if 'gevent.monkey' in sys.modules:
 else:
    from six.moves.queue import Queue, Empty  # noqa

+import six
 from six.moves import range

 from cassandra import ConsistencyLevel, AuthenticationFailed, OperationTimedOut
@@ -33,24 +34,15 @@ from cassandra.protocol import (ReadyMessage, AuthenticateMessage, OptionsMessag
                                StartupMessage, ErrorMessage, CredentialsMessage,
                                QueryMessage, ResultMessage, decode_response,
                                InvalidRequestException, SupportedMessage)
-import six
+from cassandra.util import OrderedDict


 log = logging.getLogger(__name__)

-locally_supported_compressions = {}
-
-try:
-    import snappy
-except ImportError:
-    pass
-else:
-    # work around apparently buggy snappy decompress
-    def decompress(byts):
-        if byts == '\x00':
-            return ''
-        return snappy.decompress(byts)
-    locally_supported_compressions['snappy'] = (snappy.compress, decompress)
+# We use an ordered dictionary and specifically add lz4 before
+# snappy so that lz4 will be preferred. Changing the order of this
+# will change the compression preferences for the driver.
+locally_supported_compressions = OrderedDict()

 try:
    import lz4
@@ -72,6 +64,18 @@ else:

    locally_supported_compressions['lz4'] = (lz4_compress, lz4_decompress)

+try:
+    import snappy
+except ImportError:
+    pass
+else:
+    # work around apparently buggy snappy decompress
+    def decompress(byts):
+        if byts == '\x00':
+            return ''
+        return snappy.decompress(byts)
+    locally_supported_compressions['snappy'] = (snappy.compress, decompress)
+

 MAX_STREAM_PER_CONNECTION = 127

@@ -380,7 +384,22 @@ class Connection(object):
                          locally_supported_compressions.keys(),
                          remote_supported_compressions)
            else:
-                compression_type = next(iter(overlap))  # choose any
+                compression_type = None
+                if isinstance(self.compression, basestring):
+                    # the user picked a specific compression type ('snappy' or 'lz4')
+                    if self.compression not in remote_supported_compressions:
+                        raise ProtocolError(
+                            "The requested compression type (%s) is not supported by the Cassandra server at %s"
+                            % (self.compression, self.host))
+                    compression_type = self.compression
+                else:
+                    # our locally supported compressions are ordered to prefer
+                    # lz4, if available
+                    for k in locally_supported_compressions.keys():
+                        if k in overlap:
+                            compression_type = k
+                            break
+
                # set the decompressor here, but set the compressor only after
                # a successful Ready message
                self._compressor, self.decompressor = \
--- a/cassandra/metrics.py
+++ b/cassandra/metrics.py
@@ -15,7 +15,12 @@
 from itertools import chain
 import logging

-from greplin import scales
+try:
+    from greplin import scales
+except ImportError:
+    raise ImportError(
+        "The scales library is required for metrics support: "
+        "https://pypi.python.org/pypi/scales")

 log = logging.getLogger(__name__)

--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1 @@
-blist
 futures
-scales >=1.0.6
-six >=1.6
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,6 @@ from distutils.errors import (CCompilerError, DistutilsPlatformError,
 from distutils.cmd import Command


-import platform
 import os
 import warnings

@@ -181,9 +180,7 @@ def run_setup(extensions):
        kw['cmdclass']['build_ext'] = build_extensions
        kw['ext_modules'] = extensions

-    dependencies = ['futures', 'scales >=1.0.6', 'blist', 'six >=1.6']
-    if platform.python_implementation() != "CPython":
-        dependencies.remove('blist')
+    dependencies = ['futures', 'six >=1.6']

    setup(
        name='cassandra-driver',
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -1,3 +1,5 @@
+blist
+scales
 nose
 mock
 ccm
--- a/tests/unit/test_connection.py
+++ b/tests/unit/test_connection.py
@@ -24,10 +24,11 @@ from mock import Mock, ANY

 from cassandra.cluster import Cluster
 from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT,
-                                  HEADER_DIRECTION_FROM_CLIENT, ProtocolError)
+                                  HEADER_DIRECTION_FROM_CLIENT, ProtocolError,
+                                  locally_supported_compressions)
+from cassandra.marshal import uint8_pack, uint32_pack
 from cassandra.protocol import (write_stringmultimap, write_int, write_string,
                                SupportedMessage)
-from cassandra.marshal import uint8_pack, uint32_pack


 class ConnectionTest(unittest.TestCase):
@@ -146,6 +147,123 @@ class ConnectionTest(unittest.TestCase):
        args, kwargs = c.defunct.call_args
        self.assertIsInstance(args[0], ProtocolError)

+    def test_prefer_lz4_compression(self, *args):
+        c = self.make_connection()
+        c._id_queue.get_nowait()
+        c._callbacks = {0: c._handle_options_response}
+        c.defunct = Mock()
+        c.cql_version = "3.0.3"
+
+        locally_supported_compressions.pop('lz4', None)
+        locally_supported_compressions.pop('snappy', None)
+        locally_supported_compressions['lz4'] = ('lz4compress', 'lz4decompress')
+        locally_supported_compressions['snappy'] = ('snappycompress', 'snappydecompress')
+
+        # read in a SupportedMessage response
+        header = self.make_header_prefix(SupportedMessage)
+
+        options_buf = BytesIO()
+        write_stringmultimap(options_buf, {
+            'CQL_VERSION': ['3.0.3'],
+            'COMPRESSION': ['snappy', 'lz4']
+        })
+        options = options_buf.getvalue()
+
+        message = self.make_msg(header, options)
+        c.process_msg(message, len(message) - 8)
+
+        self.assertEqual(c.decompressor, locally_supported_compressions['lz4'][1])
+
+    def test_requested_compression_not_available(self, *args):
+        c = self.make_connection()
+        c._id_queue.get_nowait()
+        c._callbacks = {0: c._handle_options_response}
+        c.defunct = Mock()
+        # request lz4 compression
+        c.compression = "lz4"
+
+        locally_supported_compressions.pop('lz4', None)
+        locally_supported_compressions.pop('snappy', None)
+        locally_supported_compressions['lz4'] = ('lz4compress', 'lz4decompress')
+        locally_supported_compressions['snappy'] = ('snappycompress', 'snappydecompress')
+
+        # read in a SupportedMessage response
+        header = self.make_header_prefix(SupportedMessage)
+
+        # the server only supports snappy
+        options_buf = BytesIO()
+        write_stringmultimap(options_buf, {
+            'CQL_VERSION': ['3.0.3'],
+            'COMPRESSION': ['snappy']
+        })
+        options = options_buf.getvalue()
+
+        message = self.make_msg(header, options)
+        c.process_msg(message, len(message) - 8)
+
+        # make sure it errored correctly
+        c.defunct.assert_called_once_with(ANY)
+        args, kwargs = c.defunct.call_args
+        self.assertIsInstance(args[0], ProtocolError)
+
+    def test_use_requested_compression(self, *args):
+        c = self.make_connection()
+        c._id_queue.get_nowait()
+        c._callbacks = {0: c._handle_options_response}
+        c.defunct = Mock()
+        # request snappy compression
+        c.compression = "snappy"
+
+        locally_supported_compressions.pop('lz4', None)
+        locally_supported_compressions.pop('snappy', None)
+        locally_supported_compressions['lz4'] = ('lz4compress', 'lz4decompress')
+        locally_supported_compressions['snappy'] = ('snappycompress', 'snappydecompress')
+
+        # read in a SupportedMessage response
+        header = self.make_header_prefix(SupportedMessage)
+
+        # the server only supports snappy
+        options_buf = BytesIO()
+        write_stringmultimap(options_buf, {
+            'CQL_VERSION': ['3.0.3'],
+            'COMPRESSION': ['snappy', 'lz4']
+        })
+        options = options_buf.getvalue()
+
+        message = self.make_msg(header, options)
+        c.process_msg(message, len(message) - 8)
+
+        self.assertEqual(c.decompressor, locally_supported_compressions['snappy'][1])
+
+    def test_disable_compression(self, *args):
+        c = self.make_connection()
+        c._id_queue.get_nowait()
+        c._callbacks = {0: c._handle_options_response}
+        c.defunct = Mock()
+        # disable compression
+        c.compression = False
+
+        locally_supported_compressions.pop('lz4', None)
+        locally_supported_compressions.pop('snappy', None)
+        locally_supported_compressions['lz4'] = ('lz4compress', 'lz4decompress')
+        locally_supported_compressions['snappy'] = ('snappycompress', 'snappydecompress')
+
+        # read in a SupportedMessage response
+        header = self.make_header_prefix(SupportedMessage)
+
+        # the server only supports snappy
+        options_buf = BytesIO()
+        write_stringmultimap(options_buf, {
+            'CQL_VERSION': ['3.0.3'],
+            'COMPRESSION': ['snappy', 'lz4']
+        })
+        options = options_buf.getvalue()
+
+        message = self.make_msg(header, options)
+        c.process_msg(message, len(message) - 8)
+
+        self.assertEqual(c.decompressor, None)
+
    def test_not_implemented(self):
        """
        Ensure the following methods throw NIE's. If not, come back and test them.
--- a/tox.ini
+++ b/tox.ini
@@ -18,3 +18,19 @@ deps = nose
       pip
       PyYAML
       six
+       scales
+       blist
+commands = {envpython} setup.py build_ext --inplace
+           nosetests --verbosity=2 tests/unit/
+
+[testenv:pypy]
+deps = nose
+       mock
+       ccm
+       unittest2
+       pip
+       PyYAML
+       scales
+       six
+commands = {envpython} setup.py build_ext --inplace
+           nosetests --verbosity=2 tests/unit/