Add support for TCompactProtocol (#70)

2 files changed, 28 insertions, 3 deletions

diff --git a/NEWS.rst b/NEWS.rst
index bd6eb13..5875e50 100644
--- a/NEWS.rst
+++ b/NEWS.rst
@@ -9,6 +9,9 @@ HappyBase 0.9
 
 Release date: *not yet released*
 
+* Add support for the Thrift compact protocol (``TCompactProtocol``) in
+  :py:class:`Connection`.
+
 
 HappyBase 0.8
 -------------
diff --git a/happybase/connection.py b/happybase/connection.py
index 11f353e..e52d256 100644
--- a/happybase/connection.py
+++ b/happybase/connection.py
@@ -8,7 +8,7 @@ import logging
 
 from thrift.transport.TSocket import TSocket
 from thrift.transport.TTransport import TBufferedTransport, TFramedTransport
-from thrift.protocol import TBinaryProtocol
+from thrift.protocol import TBinaryProtocol, TCompactProtocol
 
 from .hbase import Hbase
 from .hbase.ttypes import ColumnDescriptor
@@ -22,11 +22,16 @@ THRIFT_TRANSPORTS = dict(
     buffered=TBufferedTransport,
     framed=TFramedTransport,
 )
+THRIFT_PROTOCOLS = dict(
+    binary=TBinaryProtocol.TBinaryProtocolAccelerated,
+    compact=TCompactProtocol.TCompactProtocol,
+)
 
 DEFAULT_HOST = 'localhost'
 DEFAULT_PORT = 9090
 DEFAULT_TRANSPORT = 'buffered'
 DEFAULT_COMPAT = '0.96'
+DEFAULT_PROTOCOL = 'binary'
 
 
 class Connection(object):
@@ -67,6 +72,18 @@ class Connection(object):
     ``-hsha``, ``-nonblocking``, and ``-threadedselector`` modes use the
     framed transport.
 
+    The optional `protocol` argument specifies the Thrift transport
+    protocol to use. Supported values for this argument are ``binary``
+    (the default) and ``compact``. Make sure to choose the right one,
+    since otherwise you might see non-obvious connection errors or
+    program hangs when making a connection. ``TCompactProtocol`` is
+    a more compact binary format that is  typically more efficient to
+    process as well. ``TBinaryAccelerated`` is the default protocol that
+    happybase uses.
+
+    .. versionadded:: 0.9
+       `protocol` argument
+
     .. versionadded:: 0.5
        `timeout` argument
 
@@ -88,7 +105,7 @@ class Connection(object):
     def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
                  autoconnect=True, table_prefix=None,
                  table_prefix_separator='_', compat=DEFAULT_COMPAT,
-                 transport=DEFAULT_TRANSPORT):
+                 transport=DEFAULT_TRANSPORT, protocol=DEFAULT_PROTOCOL):
 
         if transport not in THRIFT_TRANSPORTS:
             raise ValueError("'transport' must be one of %s"
@@ -105,6 +122,10 @@ class Connection(object):
             raise ValueError("'compat' must be one of %s"
                              % ", ".join(COMPAT_MODES))
 
+        if protocol not in THRIFT_PROTOCOLS:
+            raise ValueError("'protocol' must be one of %s"
+                             % ", ".join(THRIFT_PROTOCOLS))
+
         # Allow host and port to be None, which may be easier for
         # applications wrapping a Connection instance.
         self.host = host or DEFAULT_HOST
@@ -115,6 +136,7 @@ class Connection(object):
         self.compat = compat
 
         self._transport_class = THRIFT_TRANSPORTS[transport]
+        self._protocol_class = THRIFT_PROTOCOLS[protocol]
         self._refresh_thrift_client()
 
         if autoconnect:
@@ -129,7 +151,7 @@ class Connection(object):
             socket.setTimeout(self.timeout)
 
         self.transport = self._transport_class(socket)
-        protocol = TBinaryProtocol.TBinaryProtocolAccelerated(self.transport)
+        protocol = self._protocol_class(self.transport)
         self.client = Hbase.Client(protocol)
 
     def _table_name(self, name):