Thrift SSL support

elasticsearch/connection/__init__.py

@@ -1,4 +1,4 @@
 from .base import Connection
 from .http import RequestsHttpConnection, Urllib3HttpConnection
 from .memcached import MemcachedConnection
-from .thrift import ThriftConnection
+from .thrift import ThriftConnection, THRIFT_AVAILABLE

elasticsearch/connection/thrift.py

@@ -5,7 +5,7 @@ try:
     from .esthrift import Rest
     from .esthrift.ttypes import Method, RestRequest
 
-    from thrift.transport import TTransport, TSocket
+    from thrift.transport import TTransport, TSocket, TSSLSocket
     from thrift.protocol import TBinaryProtocol
     from thrift.Thrift import TException
     THRIFT_AVAILABLE = True
@@ -23,7 +23,7 @@ class ThriftConnection(PoolingConnection):
     """
     transport_schema = 'thrift'
 
-    def __init__(self, host='localhost', port=9500, framed_transport=False, **kwargs):
+    def __init__(self, host='localhost', port=9500, framed_transport=False, use_ssl=False, **kwargs):
         """
         :arg framed_transport: use `TTransport.TFramedTransport` instead of
             `TTransport.TBufferedTransport`
@@ -33,10 +33,13 @@ class ThriftConnection(PoolingConnection):
 
         super(ThriftConnection, self).__init__(host=host, port=port, **kwargs)
         self._framed_transport = framed_transport
+        self._tsocket_class = TSocket.TSocket
+        if use_ssl:
+            self._tsocket_class = TSSLSocket.TSSLSocket 
         self._tsocket_args = (host, port)
 
     def _make_connection(self):
-        socket = TSocket.TSocket(*self._tsocket_args)
+        socket = self._tsocket_class(*self._tsocket_args)
         socket.setTimeout(self.timeout * 1000.0)
         if self._framed_transport:
             transport = TTransport.TFramedTransport(socket)

test_elasticsearch/test_connection.py

@@ -3,9 +3,27 @@ from mock import Mock, patch
 import urllib3
 
 from elasticsearch.exceptions import TransportError
-from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection
+from elasticsearch.connection import RequestsHttpConnection, \
+    Urllib3HttpConnection, THRIFT_AVAILABLE, ThriftConnection
+
+from .test_cases import TestCase, SkipTest
+
+class TestThriftConnection(TestCase):
+    def setUp(self):
+        if not THRIFT_AVAILABLE:
+            raise SkipTest('Thrift is not available.')
+        super(TestThriftConnection, self).setUp()
+
+    def test_use_ssl_uses_ssl_socket(self):
+        from thrift.transport import TSSLSocket
+        con = ThriftConnection(use_ssl=True)
+        self.assertIs(con._tsocket_class, TSSLSocket.TSSLSocket)
+
+    def test_use_normal_tsocket_by_default(self):
+        from thrift.transport import TSocket
+        con = ThriftConnection()
+        self.assertIs(con._tsocket_class, TSocket.TSocket)
 
-from .test_cases import TestCase
 
 class TestUrllib3Connection(TestCase):
     def test_http_auth(self):