dropping sasl dependency, hooray pure-sasl!

pyhs2/connections.py

@@ -1,78 +1,47 @@
-import sys
-
-from thrift.protocol.TBinaryProtocol import TBinaryProtocol
-from thrift.transport.TSocket import TSocket
 from thrift.transport.TTransport import TBufferedTransport
-import sasl
-from cloudera.thrift_sasl import TSaslClientTransport
-
+from TCLIService.ttypes import TOpenSessionReq, TCloseSessionReq
+from thrift import Thrift
+from thrift.transport import TSocket, TTransport
+from thrift.protocol.TBinaryProtocol import TBinaryProtocol
 from TCLIService import TCLIService
-
 from cursor import Cursor
-from TCLIService.ttypes import TCloseSessionReq,TOpenSessionReq
+from twitter.thrift_sasl import TSaslClientTransport
 
-class Connection(object):
-    DEFAULT_KRB_SERVICE = 'hive'
-    client = None
-    session = None
 
-    def __init__(self, host=None, port=10000, authMechanism=None, user=None, password=None, database=None, configuration=None, timeout=None):
-        authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
+class Connection(object):
+    def __init__(self, host=None, port=10000, authMechanism="PLAIN", username=None, password=None, database=None,
+                 configuration=None, timeout=None):
+        authMechanisms = {"PLAIN", "NOSASL"}
         if authMechanism not in authMechanisms:
-            raise NotImplementedError('authMechanism is either not supported or not implemented')
-        #Must set a password for thrift, even if it doesn't need one
-        #Open issue with python-sasl
-        if authMechanism == 'PLAIN' and (password is None or len(password) == 0):
-            password = 'password'
-        socket = TSocket(host, port)
-        socket.setTimeout(timeout)
-        if authMechanism == 'NOSASL':
-            transport = TBufferedTransport(socket)
-        else:
-            sasl_mech = 'PLAIN'
-            saslc = sasl.Client()
-            saslc.setAttr("username", user)
-            saslc.setAttr("password", password)
-            if authMechanism == 'KERBEROS':
-                krb_host,krb_service = self._get_krb_settings(host, configuration)
-                sasl_mech = 'GSSAPI'
-                saslc.setAttr("host", krb_host)
-                saslc.setAttr("service", krb_service)
+            raise NotImplementedError("authMechanism '{}' is either not supported or not implemented".format(authMechanism))
 
-            saslc.init()
-            transport = TSaslClientTransport(saslc, sasl_mech, socket)
+        socket = TSocket.TSocket(host, port)
+        socket.setTimeout(timeout)
 
+        if authMechanism == "NOSASL":
+            transport = TBufferedTransport(socket)
+        else:  # authMechanism == "PLAIN":
+            password = "password" if (password is None or len(password) == 0) else password
+            transport = TSaslClientTransport(socket, host=host, service=None, mechanism=authMechanism,
+                                             username=username, password=password)
         self.client = TCLIService.Client(TBinaryProtocol(transport))
         transport.open()
-        res = self.client.OpenSession(TOpenSessionReq(username=user, password=password, configuration=configuration))
+        res = self.client.OpenSession(TOpenSessionReq(username=username, password=password, configuration=configuration))
         self.session = res.sessionHandle
         if database is not None:
             with self.cursor() as cur:
                 query = "USE {0}".format(database)
-                cur.execute(query) 
+                cur.execute(query)
 
     def __enter__(self):
         return self
 
     def __exit__(self, _exc_type, _exc_value, _traceback):
         self.close()
 
-    def _get_krb_settings(self, default_host, config):
-        host = default_host
-        service = self.DEFAULT_KRB_SERVICE
-
-        if config is not None:
-            if 'krb_host' in config:
-                host = config['krb_host']
-
-            if 'krb_service' in config:
-                service = config['krb_service']
-
-        return host, service
-
     def cursor(self):
         return Cursor(self.client, self.session)
 
     def close(self):
         req = TCloseSessionReq(sessionHandle=self.session)
-        self.client.CloseSession(req)
+        self.client.CloseSession(req)

pyhs2/twitter/__init__.py

@@ -0,0 +1 @@
+__all__ = ["tsasclienttransport"]