Add pyhive connection timeout

python/.gitignore

@@ -16,4 +16,5 @@ cover/
 *.iml
 /scripts/.thrift_gen
 venv/
+.venv/
 .envrc

python/pyhive/hive.py

@@ -160,7 +160,8 @@ def __init__(
         check_hostname=None,
         ssl_cert=None,
         thrift_transport=None,
-        ssl_context=None
+        ssl_context=None,
+        connection_timeout=None,
     ):
         """Connect to HiveServer2
 
@@ -175,6 +176,7 @@ def __init__(
             Incompatible with host, port, auth, kerberos_service_name, and password.
         :param ssl_context: A custom SSL context to use for HTTPS connections. If provided,
             this overrides check_hostname and ssl_cert parameters.
+        :param connection_timeout: Millisecond timeout for Thrift connections. Skipped if using thrift_transport.
         The way to support LDAP and GSSAPI is originated from cloudera/Impyla:
         https://github.com/cloudera/impyla/blob/255b07ed973d47a3395214ed92d35ec0615ebf62
         /impala/_thrift_api.py#L152-L160
@@ -193,6 +195,8 @@ def __init__(
                 ),
                 ssl_context=ssl_context,
             )
+            if connection_timeout is not None:
+                thrift_transport.setTimeout(connection_timeout)
 
             if auth in ("BASIC", "NOSASL", "NONE", None):
                 # Always needs the Authorization header
@@ -236,6 +240,8 @@ def __init__(
             if auth is None:
                 auth = 'NONE'
             socket = thrift.transport.TSocket.TSocket(host, port)
+            if connection_timeout is not None:
+                socket.setTimeout(connection_timeout)
             if auth == 'NOSASL':
                 # NOSASL corresponds to hive.server2.authentication=NOSASL in hive-site.xml
                 self._transport = thrift.transport.TTransport.TBufferedTransport(socket)

python/pyhive/tests/test_hive.py

@@ -259,6 +259,37 @@ def test_basic_ssl_context(self):
                 cursor.execute('SELECT 1 FROM one_row')
                 self.assertEqual(cursor.fetchall(), [(1,)])
 
+    def test_connection_timeout_sasl(self):
+        """Test that a connection timeout is set without error."""
+        with contextlib.closing(hive.connect(
+            host=_HOST,
+            port=10000,
+            connection_timeout=2_000,
+        )) as connection:
+            # thrift converts milliseconds to seconds
+            assert connection._transport._trans._timeout == 2
+
+    def test_connection_timeout_nosasl(self):
+        """Test that a connection timeout is set without error."""
+        with contextlib.closing(hive.connect(
+            host=_HOST,
+            port=10000,
+            connection_timeout=2_000,
+            auth='NOSASL',
+        )) as connection:
+            # thrift converts milliseconds to seconds
+            assert connection._transport._TBufferedTransport__trans._timeout == 2
+
+    def test_connection_timeout_http(self):
+        """Test that a connection timeout is set without error."""
+        with contextlib.closing(hive.connect(
+            host=_HOST,
+            port=10000,
+            connection_timeout=2_000,
+            scheme='http',
+        )) as connection:
+            # thrift converts milliseconds to seconds
+            assert connection._transport._THttpClient__timeout == 2
 
 def _restart_hs2():
     subprocess.check_call(['sudo', 'service', 'hive-server2', 'restart'])