diff --git a/impala/dbapi.py b/impala/dbapi.py index c7b37b0f9..4249ca2d5 100644 --- a/impala/dbapi.py +++ b/impala/dbapi.py @@ -44,7 +44,7 @@ def connect(host='localhost', port=21050, database=None, timeout=None, ldap_user=None, ldap_password=None, use_kerberos=None, protocol=None, krb_host=None, use_http_transport=False, http_path='', auth_cookie_names=None, http_cookie_names=None, - retries=3, jwt=None, user_agent=None): + retries=3, jwt=None, user_agent=None, cursor_configuration=None): """Get a connection to HiveServer2 (HS2). These options are largely compatible with the impala-shell command line @@ -103,6 +103,8 @@ def connect(host='localhost', port=21050, database=None, timeout=None, This is used for auth_mechanism=JWT when using the HTTP transport. user_agent: A user specified user agent when HTTP transport is used. If none is specified, 'Python/ImpylaHttpClient' is used + cursor_configuration: object that contains configuration options that will be passed to + each cursor, optional use_ldap : bool, optional Specify `auth_mechanism='LDAP'` instead. @@ -201,7 +203,7 @@ def connect(host='localhost', port=21050, database=None, timeout=None, http_cookie_names=http_cookie_names, retries=retries, jwt=jwt, user_agent=user_agent) - return hs2.HiveServer2Connection(service, default_db=database) + return hs2.HiveServer2Connection(service, default_db=database, cursor_configuration=cursor_configuration) class _DBAPITypeObject(object): diff --git a/impala/hiveserver2.py b/impala/hiveserver2.py index 73f4555ff..fd0c6dc95 100644 --- a/impala/hiveserver2.py +++ b/impala/hiveserver2.py @@ -58,11 +58,12 @@ class HiveServer2Connection(Connection): # thrift service # it's instantiated with an alive TCLIService.Client - def __init__(self, service, default_db=None): + def __init__(self, service, default_db=None, cursor_configuration=None): log.debug('HiveServer2Connection(service=%s, default_db=%s)', service, default_db) self.service = service self.default_db = default_db + self.cursor_configuration = cursor_configuration def close(self): """Close the session and the Thrift transport.""" @@ -126,7 +127,11 @@ def cursor(self, user=None, configuration=None, convert_types=True, log.debug('.cursor(): getting new session_handle') - session = self.service.open_session(user, configuration) + session_configuration = {} + for conf in [self.cursor_configuration, configuration]: + if conf is not None: + session_configuration.update(conf) + session = self.service.open_session(user, session_configuration if session_configuration else None) log.debug('HiveServer2Cursor(service=%s, session_handle=%s, ' 'default_config=%s, hs2_protocol_version=%s)',