Jun 7, 2023 · Jun 6, 2023 · Jun 6, 2023 · Jun 6, 2023 · Jun 6, 2023 · Jun 6, 2023
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
 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)
        # _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"):
            ssl_context=ssl_context,
        )

        timeout = kwargs.get("_socket_timeout")
        # setTimeout defaults toNone (i.e. no timeout), and is expected in ms
        timeout = kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT)
        # setTimeout defaults to15 minutes and is expected in ms
        self._transport.setTimeout(timeout and (float(timeout) * 1000.0))

        self._transport.setCustomHeaders(dict(http_headers))
diff --git a/tests/e2e/driver_tests.py b/tests/e2e/driver_tests.py
 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

    @skipIf(pysql_has_version('<', '2'), 'requires pysql v2')
    def test_socket_timeout(self):
        #  Weweexpect 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:
diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py
        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)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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)
Expand DownExpand Up		@@ -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"):
Expand DownExpand Up		@@ -152,8 +153,8 @@ def __init__(
		ssl_context=ssl_context,
		)

		timeout = kwargs.get("_socket_timeout")
		# setTimeout defaults toNone (i.e. no timeout), and is expected in ms
		timeout = kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT)
		# setTimeout defaults to15 minutes and is expected in ms
		self._transport.setTimeout(timeout and (float(timeout) * 1000.0))

		self._transport.setCustomHeaders(dict(http_headers))
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand DownExpand Up		@@ -459,14 +459,25 @@ def test_temp_view_fetch(self):

		@skipIf(pysql_has_version('<', '2'), 'requires pysql v2')
		def test_socket_timeout(self):
		# Weweexpect 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:
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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)

Expand Down