diff --git a/pyhive/hive.py b/pyhive/hive.py index 3f71df33..d0701481 100644 --- a/pyhive/hive.py +++ b/pyhive/hive.py @@ -120,7 +120,8 @@ def __init__( password=None, check_hostname=None, ssl_cert=None, - thrift_transport=None + thrift_transport=None, + timeout=None ): """Connect to HiveServer2 @@ -152,6 +153,7 @@ def __init__( ), ssl_context=ssl_context, ) + thrift_transport.setTimeout(timeout) if auth in ("BASIC", "NOSASL", "NONE", None): # Always needs the Authorization header @@ -195,6 +197,7 @@ def __init__( if auth is None: auth = 'NONE' socket = thrift.transport.TSocket.TSocket(host, port) + socket.setTimeout(timeout) if auth == 'NOSASL': # NOSASL corresponds to hive.server2.authentication=NOSASL in hive-site.xml self._transport = thrift.transport.TTransport.TBufferedTransport(socket)