diff --git a/elasticsearch/connection/__init__.py b/elasticsearch/connection/__init__.py index 6f2ea069..daf4ed15 100644 --- a/elasticsearch/connection/__init__.py +++ b/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 diff --git a/elasticsearch/connection/thrift.py b/elasticsearch/connection/thrift.py index d7c9b9a3..85e64312 100644 --- a/elasticsearch/connection/thrift.py +++ b/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) diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index 13f0fa68..31ec10d2 100644 --- a/test_elasticsearch/test_connection.py +++ b/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):