Add socket timeout param

Author: NiallEgan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -60,6 +60,9 @@ def __init__(self,
         #   Set client SSL certificate.
         # _retry_stop_after_attempts_count
         #  The maximum number of attempts during a request retry sequence (defaults to 24)
+        # _socket_timeout
+        #  The timeout in seconds for socket send, recv and connect operations. Defaults to None for
+        #  no timeout. Should be a positive float or integer.
 
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)

cmdexec/clients/python/src/databricks/sql/thrift_backend.py

@@ -71,6 +71,11 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
         #   next calculated pre-retry delay would go past
         #   _retry_stop_after_attempts_duration, stop now.)
         #
+        # _retry_stop_after_attempts_count
+        #  The maximum number of times we should retry retryable requests (defaults to 24)
+        # _socket_timeout
+        #  The timeout in seconds for socket send, recv and connect operations. Defaults to None for
+        #  no timeout. Should be a positive float or integer.
 
         port = port or 443
         if kwargs.get("_connection_uri"):
@@ -109,6 +114,10 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
             ssl_context=ssl_context,
         )
 
+        timeout = kwargs.get("_socket_timeout")
+        # setTimeout defaults to None (i.e. no timeout), and is expected in ms
+        self._transport.setTimeout(timeout and (float(timeout) * 1000.0))
+
         self._transport.setCustomHeaders(dict(http_headers))
         protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(self._transport)
         self._client = TCLIService.Client(protocol)

cmdexec/clients/python/tests/test_thrift_backend.py

@@ -193,6 +193,15 @@ def test_port_and_host_are_respected(self, t_http_client_class):
         self.assertEqual(t_http_client_class.call_args[1]["uri_or_host"],
                          "https://hostname:123/path_value")
 
+    @patch("thrift.transport.THttpClient.THttpClient")
+    def test_socket_timeout_is_propagated(self, t_http_client_class):
+        ThriftBackend("hostname", 123, "path_value", [], _socket_timeout=129)
+        self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], 129 * 1000)
+        ThriftBackend("hostname", 123, "path_value", [], _socket_timeout=0)
+        self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], 0)
+        ThriftBackend("hostname", 123, "path_value", [], _socket_timeout=None)
+        self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], None)
+
     def test_non_primitive_types_raise_error(self):
         columns = [
             ttypes.TColumnDesc(

cmdexec/clients/python/tests/tests.py

@@ -333,6 +333,11 @@ def test_max_number_of_retries_passthrough(self, mock_client_class):
 
         self.assertEqual(mock_client_class.call_args[1]["_retry_stop_after_attempts_count"], 54)
 
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_socket_timeout_passthrough(self, mock_client_class):
+        databricks.sql.connect(_socket_timeout=234, **self.DUMMY_CONNECTION_ARGS)
+        self.assertEqual(mock_client_class.call_args[1]["_socket_timeout"], 234)
+
     def test_version_is_canonical(self):
         version = databricks.sql.__version__
         canonical_version_re = r'^([1-9][0-9]*!)?(0|[1-9][0-9]*)(\.(0|[1-9][0-9]*))*((a|b|rc)' \