def get_thrift_transport()

in src/ab/plugins/db/hive.py [0:0]
40 lines of code
8 McCabe index (conditional complexity)

def get_thrift_transport(host=None, port=None, username=None, auth=None, kerberos_service_name=None, password=None,
                         kerberos_host=None, socket_timeout_in_ms=None):
    """copy from pyhive/hive.py and patch"""
    if port is None:
        port = 10000
    if auth is None:
        auth = 'NONE'
    socket = thrift.transport.TSocket.TSocket(host, port)

    #################################################
    #             1. set socket timeout             #
    #################################################
    if socket_timeout_in_ms is not None:
        socket.setTimeout(socket_timeout_in_ms)  # in ms

    if auth == 'NOSASL':
        # NOSASL corresponds to hive.server2.authentication=NOSASL in hive-site.xml
        return thrift.transport.TTransport.TBufferedTransport(socket)
    elif auth in ('LDAP', 'KERBEROS', 'NONE', 'CUSTOM'):
        # Defer import so package dependency is optional
        import sasl
        import thrift_sasl

        if auth == 'KERBEROS':
            # KERBEROS mode in hive.server2.authentication is GSSAPI in sasl library
            sasl_auth = 'GSSAPI'
        else:
            sasl_auth = 'PLAIN'
            if password is None:
                # Password doesn't matter in NONE mode, just needs to be nonempty.
                password = 'x'

        def sasl_factory():
            sasl_client = sasl.Client()
            sasl_client.setAttr('host', host)
            if sasl_auth == 'GSSAPI':
                #################################################
                #          2. adapt for huawei hadoop           #
                #################################################
                if kerberos_host is not None:
                    sasl_client.setAttr('host', kerberos_host)

                sasl_client.setAttr('service', kerberos_service_name)
            elif sasl_auth == 'PLAIN':
                sasl_client.setAttr('username', username)
                sasl_client.setAttr('password', password)
            else:
                raise AssertionError

            #################################################
            #         3. copy from huawei just in case      #
            #################################################
            sasl_client.setAttr("maxbufsize", 16777215)

            sasl_client.init()
            return sasl_client

        return thrift_sasl.TSaslClientTransport(sasl_factory, sasl_auth, socket)
    else:
        # All HS2 config options:
        # https://cwiki.apache.org/confluence/display/Hive/Setting+Up+HiveServer2#SettingUpHiveServer2-Configuration
        # PAM currently left to end user via thrift_transport option.
        raise NotImplementedError(
            "Only NONE, NOSASL, LDAP, KERBEROS, CUSTOM "
            "authentication are supported, got {}".format(auth))