Default socket timeout to 15 min (#137)

Signed-off-by: Matthew Kim <[email protected]>

Author: mattdeekay

src/databricks/sql/thrift_backend.py

@@ -37,6 +37,7 @@
 DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase"
 
 TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString"
+DEFAULT_SOCKET_TIMEOUT = float(900)
 
 # see Connection.__init__ for parameter descriptions.
 # - Min/Max avoids unsustainable configs (sane values are far more constrained)
@@ -99,8 +100,8 @@ def __init__(
         # _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.
+        #  The timeout in seconds for socket send, recv and connect operations. Should be a positive float or integer.
+        #  (defaults to 900)
 
         port = port or 443
         if kwargs.get("_connection_uri"):
@@ -152,8 +153,8 @@ def __init__(
             ssl_context=ssl_context,
         )
 
-        timeout = kwargs.get("_socket_timeout")
-        # setTimeout defaults to None (i.e. no timeout), and is expected in ms
+        timeout = kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT)
+        # setTimeout defaults to 15 minutes and is expected in ms
         self._transport.setTimeout(timeout and (float(timeout) * 1000.0))
 
         self._transport.setCustomHeaders(dict(http_headers))

tests/e2e/driver_tests.py

@@ -18,7 +18,7 @@
 import pytest
 
 import databricks.sql as sql
-from databricks.sql import STRING, BINARY, NUMBER, DATETIME, DATE, DatabaseError, Error, OperationalError
+from databricks.sql import STRING, BINARY, NUMBER, DATETIME, DATE, DatabaseError, Error, OperationalError, RequestError
 from tests.e2e.common.predicates import pysql_has_version, pysql_supports_arrow, compare_dbr_versions, is_thrift_v5_plus
 from tests.e2e.common.core_tests import CoreTestMixin, SmokeTestMixin
 from tests.e2e.common.large_queries_mixin import LargeQueriesMixin
@@ -460,14 +460,25 @@ def test_temp_view_fetch(self):
     @skipIf(pysql_has_version('<', '2'), 'requires pysql v2')
     @skipIf(True, "Unclear the purpose of this test since urllib3 does not complain when timeout == 0")
     def test_socket_timeout(self):
-        #  We we expect to see a BlockingIO error when the socket is opened
+        #  We expect to see a BlockingIO error when the socket is opened
         #  in non-blocking mode, since no poll is done before the read
         with self.assertRaises(OperationalError) as cm:
             with self.cursor({"_socket_timeout": 0}):
                 pass
 
         self.assertIsInstance(cm.exception.args[1], io.BlockingIOError)
 
+    @skipIf(pysql_has_version('<', '2'), 'requires pysql v2')
+    def test_socket_timeout_user_defined(self):
+        #  We expect to see a TimeoutError when the socket timeout is only
+        #  1 sec for a query that takes longer than that to process
+        with self.assertRaises(RequestError) as cm:
+            with self.cursor({"_socket_timeout": 1}) as cursor:
+                query = "select * from range(10000000)"
+                cursor.execute(query)
+
+        self.assertIsInstance(cm.exception.args[1], TimeoutError)
+
     def test_ssp_passthrough(self):
         for enable_ansi in (True, False):
             with self.cursor({"session_configuration": {"ansi_mode": enable_ansi}}) as cursor:

tests/unit/test_thrift_backend.py

@@ -217,6 +217,8 @@ def test_socket_timeout_is_propagated(self, t_http_client_class):
         self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], 129 * 1000)
         ThriftBackend("hostname", 123, "path_value", [], auth_provider=AuthProvider(), _socket_timeout=0)
         self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], 0)
+        ThriftBackend("hostname", 123, "path_value", [], auth_provider=AuthProvider())
+        self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], 900 * 1000)
         ThriftBackend("hostname", 123, "path_value", [], auth_provider=AuthProvider(), _socket_timeout=None)
         self.assertEqual(t_http_client_class.return_value.setTimeout.call_args[0][0], None)