From df75ba18e246caba5b0f5f8cf7902d537f0985f6 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Wed, 18 Mar 2026 10:50:12 +0530 Subject: [PATCH 1/6] Add _retry_server_directed_only mode for Retry-After header compliance When enabled, the connector only retries on 429/503 if the server includes a Retry-After header in the response. This prevents duplicate side effects for non-idempotent ExecuteStatement operations where the server has not explicitly signaled that retry is safe. The new opt-in parameter `_retry_server_directed_only` threads through ClientContext, all three DatabricksRetryPolicy construction sites (Thrift, SEA, UnifiedHttpClient), and the retry policy's should_retry/is_retry methods. Default behavior (retry without requiring the header) is unchanged. Signed-off-by: Shubham Dhal --- src/databricks/sql/auth/common.py | 2 + src/databricks/sql/auth/retry.py | 15 ++- .../sql/backend/sea/utils/http_client.py | 4 + src/databricks/sql/backend/thrift_backend.py | 4 + .../sql/common/unified_http_client.py | 1 + src/databricks/sql/utils.py | 1 + tests/unit/test_retry.py | 120 ++++++++++++++++++ 7 files changed, 145 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py index 3e0be0d2b..2b7fea865 100644 --- a/src/databricks/sql/auth/common.py +++ b/src/databricks/sql/auth/common.py @@ -47,6 +47,7 @@ def __init__( retry_stop_after_attempts_duration: Optional[float] = None, retry_delay_default: Optional[float] = None, retry_dangerous_codes: Optional[List[int]] = None, + retry_server_directed_only: Optional[bool] = None, proxy_auth_method: Optional[str] = None, pool_connections: Optional[int] = None, pool_maxsize: Optional[int] = None, @@ -79,6 +80,7 @@ def __init__( ) self.retry_delay_default = retry_delay_default or 5.0 self.retry_dangerous_codes = retry_dangerous_codes or [] + self.retry_server_directed_only = bool(retry_server_directed_only) self.proxy_auth_method = proxy_auth_method self.pool_connections = pool_connections or 10 self.pool_maxsize = pool_maxsize or 20 diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 4281883da..0091a922d 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -94,6 +94,7 @@ def __init__( stop_after_attempts_duration: float, delay_default: float, force_dangerous_codes: List[int], + server_directed_only: bool = False, urllib3_kwargs: dict = {}, ): # These values do not change from one command to the next @@ -103,6 +104,7 @@ def __init__( self.stop_after_attempts_duration = stop_after_attempts_duration self._delay_default = delay_default self.force_dangerous_codes = force_dangerous_codes + self.server_directed_only = server_directed_only # the urllib3 kwargs are a mix of configuration (some of which we override) # and counters like `total` or `connect` which may change between successive retries @@ -202,6 +204,7 @@ def new( stop_after_attempts_duration=self.stop_after_attempts_duration, delay_default=self.delay_default, force_dangerous_codes=self.force_dangerous_codes, + server_directed_only=self.server_directed_only, urllib3_kwargs={}, ) @@ -323,7 +326,9 @@ def get_backoff_time(self) -> float: return proposed_backoff - def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: + def should_retry( + self, method: str, status_code: int, has_retry_after: bool = False + ) -> Tuple[bool, str]: """This method encapsulates the connector's approach to retries. We always retry a request unless one of these conditions is met: @@ -381,6 +386,12 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: if not self._is_method_retryable(method): return False, "Only POST requests are retried" + # In server_directed_only mode, only retry when the server explicitly signals + # it's safe via a Retry-After header. This prevents duplicate side effects for + # non-idempotent operations. + if self.server_directed_only and not has_retry_after: + return (False, "server_directed_only mode: no Retry-After header present") + # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry. if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS: return ( @@ -450,7 +461,7 @@ def is_retry( Logs a debug message if the request will be retried """ - should_retry, msg = self.should_retry(method, status_code) + should_retry, msg = self.should_retry(method, status_code, has_retry_after) if should_retry: logger.debug(msg) diff --git a/src/databricks/sql/backend/sea/utils/http_client.py b/src/databricks/sql/backend/sea/utils/http_client.py index b47f2add2..8c1d5b51d 100644 --- a/src/databricks/sql/backend/sea/utils/http_client.py +++ b/src/databricks/sql/backend/sea/utils/http_client.py @@ -90,6 +90,9 @@ def __init__( ) self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) + self._retry_server_directed_only = kwargs.get( + "_retry_server_directed_only", False + ) # Connection pooling settings self.max_connections = kwargs.get("max_connections", 10) @@ -114,6 +117,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, + server_directed_only=self._retry_server_directed_only, urllib3_kwargs=urllib3_kwargs, ) else: diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index d2b10e718..fa0cd2f1f 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -189,6 +189,9 @@ def __init__( " This behaviour is deprecated and will be removed in a future release." ) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) + self._retry_server_directed_only = kwargs.get( + "_retry_server_directed_only", False + ) additional_transport_args = {} @@ -215,6 +218,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, + server_directed_only=self._retry_server_directed_only, urllib3_kwargs=urllib3_kwargs, ) diff --git a/src/databricks/sql/common/unified_http_client.py b/src/databricks/sql/common/unified_http_client.py index 7ccd69c54..69b8c2f6c 100644 --- a/src/databricks/sql/common/unified_http_client.py +++ b/src/databricks/sql/common/unified_http_client.py @@ -99,6 +99,7 @@ def _setup_pool_managers(self): stop_after_attempts_duration=self.config.retry_stop_after_attempts_duration, delay_default=self.config.retry_delay_default, force_dangerous_codes=self.config.retry_dangerous_codes, + server_directed_only=self.config.retry_server_directed_only, ) # Initialize the required attributes that DatabricksRetryPolicy expects diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 9f96e8743..748ecc515 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -919,6 +919,7 @@ def build_client_context(server_hostname: str, version: str, **kwargs): ), retry_delay_default=kwargs.get("_retry_delay_default"), retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"), + retry_server_directed_only=kwargs.get("_retry_server_directed_only"), proxy_auth_method=kwargs.get("_proxy_auth_method"), pool_connections=kwargs.get("_pool_connections"), pool_maxsize=kwargs.get("_pool_maxsize"), diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index 897a1d111..180907605 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -83,3 +83,123 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy): retry_policy.sleep(HTTPResponse(status=503)) # Internally urllib3 calls the increment function generating a new instance for every retry retry_policy = retry_policy.increment() + + @pytest.fixture() + def server_directed_retry_policy(self) -> DatabricksRetryPolicy: + return DatabricksRetryPolicy( + delay_min=1, + delay_max=30, + stop_after_attempts_count=3, + stop_after_attempts_duration=900, + delay_default=2, + force_dangerous_codes=[], + server_directed_only=True, + ) + + def test_server_directed_only__retries_with_retry_after( + self, server_directed_retry_policy + ): + """429 + Retry-After header → should retry""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.OTHER + should_retry, msg = server_directed_retry_policy.should_retry( + "POST", 429, has_retry_after=True + ) + assert should_retry is True + + def test_server_directed_only__no_retry_without_retry_after( + self, server_directed_retry_policy + ): + """429 without Retry-After header → no retry""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.OTHER + should_retry, msg = server_directed_retry_policy.should_retry( + "POST", 429, has_retry_after=False + ) + assert should_retry is False + assert "server_directed_only" in msg + + def test_server_directed_only__no_retry_503_without_header( + self, server_directed_retry_policy + ): + """503 without Retry-After header → no retry""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.OTHER + should_retry, msg = server_directed_retry_policy.should_retry( + "POST", 503, has_retry_after=False + ) + assert should_retry is False + assert "server_directed_only" in msg + + def test_server_directed_only__overrides_dangerous_codes(self): + """force_dangerous_codes=[500] + no Retry-After → no retry in server_directed_only mode""" + policy = DatabricksRetryPolicy( + delay_min=1, + delay_max=30, + stop_after_attempts_count=3, + stop_after_attempts_duration=900, + delay_default=2, + force_dangerous_codes=[500], + server_directed_only=True, + ) + policy._retry_start_time = time.time() + policy.command_type = CommandType.EXECUTE_STATEMENT + should_retry, msg = policy.should_retry("POST", 500, has_retry_after=False) + assert should_retry is False + assert "server_directed_only" in msg + + def test_server_directed_only__non_retryable_codes_unaffected( + self, server_directed_retry_policy + ): + """401/403/501 still don't retry even with Retry-After header""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.OTHER + for code in [401, 403, 501]: + should_retry, msg = server_directed_retry_policy.should_retry( + "POST", code, has_retry_after=True + ) + assert should_retry is False, f"Code {code} should never retry" + + def test_default_mode_unchanged(self, retry_policy): + """server_directed_only=False preserves existing behavior — 429 retries without header""" + retry_policy._retry_start_time = time.time() + retry_policy.command_type = CommandType.OTHER + should_retry, msg = retry_policy.should_retry( + "POST", 429, has_retry_after=False + ) + assert should_retry is True + + def test_server_directed_only__survives_new(self, server_directed_retry_policy): + """urllib3 calls .new() between retries to create a fresh policy instance. + Verify that server_directed_only is carried over and still enforced.""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.OTHER + new_policy = server_directed_retry_policy.new() + assert new_policy.server_directed_only is True + # The new instance should still block retries without Retry-After + should_retry, msg = new_policy.should_retry("POST", 429, has_retry_after=False) + assert should_retry is False + assert "server_directed_only" in msg + + def test_server_directed_only__execute_statement_with_retry_after( + self, server_directed_retry_policy + ): + """EXECUTE_STATEMENT + 429 + Retry-After header → retry""" + server_directed_retry_policy._retry_start_time = time.time() + server_directed_retry_policy.command_type = CommandType.EXECUTE_STATEMENT + should_retry, msg = server_directed_retry_policy.should_retry( + "POST", 429, has_retry_after=True + ) + assert should_retry is True + + def test_404_does_not_retry_for_any_command_type(self, retry_policy): + """Test that 404 never retries for any CommandType""" + retry_policy._retry_start_time = time.time() + + # Test for each CommandType + for command_type in CommandType: + retry_policy.command_type = command_type + should_retry, msg = retry_policy.should_retry("POST", 404) + + assert should_retry is False, f"404 should not retry for {command_type}" + assert "404" in msg or "NOT_FOUND" in msg From 2a8568866957f2adf07b43e49531416b7a4ba993 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Wed, 18 Mar 2026 15:25:57 +0530 Subject: [PATCH 2/6] Remove unnecessary _retry_server_directed_only instance variables Inline kwargs.get() at the single point of use in ThriftDatabricksClient and SeaHttpClient instead of storing as dead instance state. Signed-off-by: Shubham Dhal --- src/databricks/sql/backend/sea/utils/http_client.py | 5 +---- src/databricks/sql/backend/thrift_backend.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/databricks/sql/backend/sea/utils/http_client.py b/src/databricks/sql/backend/sea/utils/http_client.py index 8c1d5b51d..034852cbc 100644 --- a/src/databricks/sql/backend/sea/utils/http_client.py +++ b/src/databricks/sql/backend/sea/utils/http_client.py @@ -90,9 +90,6 @@ def __init__( ) self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) - self._retry_server_directed_only = kwargs.get( - "_retry_server_directed_only", False - ) # Connection pooling settings self.max_connections = kwargs.get("max_connections", 10) @@ -117,7 +114,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, - server_directed_only=self._retry_server_directed_only, + server_directed_only=kwargs.get("_retry_server_directed_only", False), urllib3_kwargs=urllib3_kwargs, ) else: diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index fa0cd2f1f..22ccc1592 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -189,9 +189,6 @@ def __init__( " This behaviour is deprecated and will be removed in a future release." ) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) - self._retry_server_directed_only = kwargs.get( - "_retry_server_directed_only", False - ) additional_transport_args = {} @@ -218,7 +215,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, - server_directed_only=self._retry_server_directed_only, + server_directed_only=kwargs.get("_retry_server_directed_only", False), urllib3_kwargs=urllib3_kwargs, ) From 8f9d84ba0e381c9d0b61379a7d4356c951049580 Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Wed, 18 Mar 2026 16:49:52 +0530 Subject: [PATCH 3/6] Address PR feedback: rename and clean up retry-after parameter - Rename server_directed_only to respect_server_retry_after_header throughout for clarity - Store _respect_server_retry_after_header as instance variable in Thrift/SEA backends to match existing kwargs extraction pattern - Replace duplicate test fixture with _make_retry_policy(**overrides) helper for flexible policy construction in tests Signed-off-by: Shubham Dhal --- src/databricks/sql/auth/common.py | 4 +- src/databricks/sql/auth/retry.py | 16 +-- .../sql/backend/sea/utils/http_client.py | 5 +- src/databricks/sql/backend/thrift_backend.py | 5 +- .../sql/common/unified_http_client.py | 2 +- src/databricks/sql/utils.py | 2 +- tests/unit/test_retry.py | 123 +++++++----------- 7 files changed, 69 insertions(+), 88 deletions(-) diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py index 2b7fea865..44151ecb3 100644 --- a/src/databricks/sql/auth/common.py +++ b/src/databricks/sql/auth/common.py @@ -47,7 +47,7 @@ def __init__( retry_stop_after_attempts_duration: Optional[float] = None, retry_delay_default: Optional[float] = None, retry_dangerous_codes: Optional[List[int]] = None, - retry_server_directed_only: Optional[bool] = None, + respect_server_retry_after_header: Optional[bool] = None, proxy_auth_method: Optional[str] = None, pool_connections: Optional[int] = None, pool_maxsize: Optional[int] = None, @@ -80,7 +80,7 @@ def __init__( ) self.retry_delay_default = retry_delay_default or 5.0 self.retry_dangerous_codes = retry_dangerous_codes or [] - self.retry_server_directed_only = bool(retry_server_directed_only) + self.respect_server_retry_after_header = bool(respect_server_retry_after_header) self.proxy_auth_method = proxy_auth_method self.pool_connections = pool_connections or 10 self.pool_maxsize = pool_maxsize or 20 diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 0091a922d..454e2057c 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -94,7 +94,7 @@ def __init__( stop_after_attempts_duration: float, delay_default: float, force_dangerous_codes: List[int], - server_directed_only: bool = False, + respect_server_retry_after_header: bool = False, urllib3_kwargs: dict = {}, ): # These values do not change from one command to the next @@ -104,7 +104,7 @@ def __init__( self.stop_after_attempts_duration = stop_after_attempts_duration self._delay_default = delay_default self.force_dangerous_codes = force_dangerous_codes - self.server_directed_only = server_directed_only + self.respect_server_retry_after_header = respect_server_retry_after_header # the urllib3 kwargs are a mix of configuration (some of which we override) # and counters like `total` or `connect` which may change between successive retries @@ -204,7 +204,7 @@ def new( stop_after_attempts_duration=self.stop_after_attempts_duration, delay_default=self.delay_default, force_dangerous_codes=self.force_dangerous_codes, - server_directed_only=self.server_directed_only, + respect_server_retry_after_header=self.respect_server_retry_after_header, urllib3_kwargs={}, ) @@ -386,11 +386,11 @@ def should_retry( if not self._is_method_retryable(method): return False, "Only POST requests are retried" - # In server_directed_only mode, only retry when the server explicitly signals - # it's safe via a Retry-After header. This prevents duplicate side effects for - # non-idempotent operations. - if self.server_directed_only and not has_retry_after: - return (False, "server_directed_only mode: no Retry-After header present") + # When respect_server_retry_after_header is enabled, only retry when the + # server explicitly signals it's safe via a Retry-After header. This prevents + # duplicate side effects for non-idempotent operations. + if self.respect_server_retry_after_header and not has_retry_after: + return (False, "respect_server_retry_after_header mode: no Retry-After header present") # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry. if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS: diff --git a/src/databricks/sql/backend/sea/utils/http_client.py b/src/databricks/sql/backend/sea/utils/http_client.py index 034852cbc..476bddb17 100644 --- a/src/databricks/sql/backend/sea/utils/http_client.py +++ b/src/databricks/sql/backend/sea/utils/http_client.py @@ -90,6 +90,9 @@ def __init__( ) self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) + self._respect_server_retry_after_header = kwargs.get( + "_respect_server_retry_after_header", False + ) # Connection pooling settings self.max_connections = kwargs.get("max_connections", 10) @@ -114,7 +117,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, - server_directed_only=kwargs.get("_retry_server_directed_only", False), + respect_server_retry_after_header=self._respect_server_retry_after_header, urllib3_kwargs=urllib3_kwargs, ) else: diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 22ccc1592..7cfe181bc 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -189,6 +189,9 @@ def __init__( " This behaviour is deprecated and will be removed in a future release." ) self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", []) + self._respect_server_retry_after_header = kwargs.get( + "_respect_server_retry_after_header", False + ) additional_transport_args = {} @@ -215,7 +218,7 @@ def __init__( stop_after_attempts_duration=self._retry_stop_after_attempts_duration, delay_default=self._retry_delay_default, force_dangerous_codes=self.force_dangerous_codes, - server_directed_only=kwargs.get("_retry_server_directed_only", False), + respect_server_retry_after_header=self._respect_server_retry_after_header, urllib3_kwargs=urllib3_kwargs, ) diff --git a/src/databricks/sql/common/unified_http_client.py b/src/databricks/sql/common/unified_http_client.py index 69b8c2f6c..72b5597cc 100644 --- a/src/databricks/sql/common/unified_http_client.py +++ b/src/databricks/sql/common/unified_http_client.py @@ -99,7 +99,7 @@ def _setup_pool_managers(self): stop_after_attempts_duration=self.config.retry_stop_after_attempts_duration, delay_default=self.config.retry_delay_default, force_dangerous_codes=self.config.retry_dangerous_codes, - server_directed_only=self.config.retry_server_directed_only, + respect_server_retry_after_header=self.config.respect_server_retry_after_header, ) # Initialize the required attributes that DatabricksRetryPolicy expects diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 748ecc515..7311c30ba 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -919,7 +919,7 @@ def build_client_context(server_hostname: str, version: str, **kwargs): ), retry_delay_default=kwargs.get("_retry_delay_default"), retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"), - retry_server_directed_only=kwargs.get("_retry_server_directed_only"), + respect_server_retry_after_header=kwargs.get("_respect_server_retry_after_header"), proxy_auth_method=kwargs.get("_proxy_auth_method"), pool_connections=kwargs.get("_pool_connections"), pool_maxsize=kwargs.get("_pool_maxsize"), diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index 180907605..7abab509d 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -7,9 +7,8 @@ class TestRetry: - @pytest.fixture() - def retry_policy(self) -> DatabricksRetryPolicy: - return DatabricksRetryPolicy( + def _make_retry_policy(self, **overrides) -> DatabricksRetryPolicy: + defaults = dict( delay_min=1, delay_max=30, stop_after_attempts_count=3, @@ -17,6 +16,12 @@ def retry_policy(self) -> DatabricksRetryPolicy: delay_default=2, force_dangerous_codes=[], ) + defaults.update(overrides) + return DatabricksRetryPolicy(**defaults) + + @pytest.fixture() + def retry_policy(self) -> DatabricksRetryPolicy: + return self._make_retry_policy() @pytest.fixture() def error_history(self) -> RequestHistory: @@ -84,84 +89,56 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy): # Internally urllib3 calls the increment function generating a new instance for every retry retry_policy = retry_policy.increment() - @pytest.fixture() - def server_directed_retry_policy(self) -> DatabricksRetryPolicy: - return DatabricksRetryPolicy( - delay_min=1, - delay_max=30, - stop_after_attempts_count=3, - stop_after_attempts_duration=900, - delay_default=2, - force_dangerous_codes=[], - server_directed_only=True, - ) - - def test_server_directed_only__retries_with_retry_after( - self, server_directed_retry_policy - ): + def test_respect_server_retry_after__retries_with_retry_after(self): """429 + Retry-After header → should retry""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.OTHER - should_retry, msg = server_directed_retry_policy.should_retry( - "POST", 429, has_retry_after=True - ) + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.OTHER + should_retry, msg = policy.should_retry("POST", 429, has_retry_after=True) assert should_retry is True - def test_server_directed_only__no_retry_without_retry_after( - self, server_directed_retry_policy - ): + def test_respect_server_retry_after__no_retry_without_retry_after(self): """429 without Retry-After header → no retry""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.OTHER - should_retry, msg = server_directed_retry_policy.should_retry( - "POST", 429, has_retry_after=False - ) + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.OTHER + should_retry, msg = policy.should_retry("POST", 429, has_retry_after=False) assert should_retry is False - assert "server_directed_only" in msg + assert "respect_server_retry_after_header" in msg - def test_server_directed_only__no_retry_503_without_header( - self, server_directed_retry_policy - ): + def test_respect_server_retry_after__no_retry_503_without_header(self): """503 without Retry-After header → no retry""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.OTHER - should_retry, msg = server_directed_retry_policy.should_retry( - "POST", 503, has_retry_after=False - ) + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.OTHER + should_retry, msg = policy.should_retry("POST", 503, has_retry_after=False) assert should_retry is False - assert "server_directed_only" in msg + assert "respect_server_retry_after_header" in msg - def test_server_directed_only__overrides_dangerous_codes(self): - """force_dangerous_codes=[500] + no Retry-After → no retry in server_directed_only mode""" - policy = DatabricksRetryPolicy( - delay_min=1, - delay_max=30, - stop_after_attempts_count=3, - stop_after_attempts_duration=900, - delay_default=2, - force_dangerous_codes=[500], - server_directed_only=True, + def test_respect_server_retry_after__overrides_dangerous_codes(self): + """force_dangerous_codes=[500] + no Retry-After → no retry in respect_server_retry_after_header mode""" + policy = self._make_retry_policy( + force_dangerous_codes=[500], respect_server_retry_after_header=True ) policy._retry_start_time = time.time() policy.command_type = CommandType.EXECUTE_STATEMENT should_retry, msg = policy.should_retry("POST", 500, has_retry_after=False) assert should_retry is False - assert "server_directed_only" in msg + assert "respect_server_retry_after_header" in msg - def test_server_directed_only__non_retryable_codes_unaffected( - self, server_directed_retry_policy - ): + def test_respect_server_retry_after__non_retryable_codes_unaffected(self): """401/403/501 still don't retry even with Retry-After header""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.OTHER + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.OTHER for code in [401, 403, 501]: - should_retry, msg = server_directed_retry_policy.should_retry( + should_retry, msg = policy.should_retry( "POST", code, has_retry_after=True ) assert should_retry is False, f"Code {code} should never retry" def test_default_mode_unchanged(self, retry_policy): - """server_directed_only=False preserves existing behavior — 429 retries without header""" + """respect_server_retry_after_header=False preserves existing behavior — 429 retries without header""" retry_policy._retry_start_time = time.time() retry_policy.command_type = CommandType.OTHER should_retry, msg = retry_policy.should_retry( @@ -169,27 +146,25 @@ def test_default_mode_unchanged(self, retry_policy): ) assert should_retry is True - def test_server_directed_only__survives_new(self, server_directed_retry_policy): + def test_respect_server_retry_after__survives_new(self): """urllib3 calls .new() between retries to create a fresh policy instance. - Verify that server_directed_only is carried over and still enforced.""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.OTHER - new_policy = server_directed_retry_policy.new() - assert new_policy.server_directed_only is True + Verify that respect_server_retry_after_header is carried over and still enforced.""" + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.OTHER + new_policy = policy.new() + assert new_policy.respect_server_retry_after_header is True # The new instance should still block retries without Retry-After should_retry, msg = new_policy.should_retry("POST", 429, has_retry_after=False) assert should_retry is False - assert "server_directed_only" in msg + assert "respect_server_retry_after_header" in msg - def test_server_directed_only__execute_statement_with_retry_after( - self, server_directed_retry_policy - ): + def test_respect_server_retry_after__execute_statement_with_retry_after(self): """EXECUTE_STATEMENT + 429 + Retry-After header → retry""" - server_directed_retry_policy._retry_start_time = time.time() - server_directed_retry_policy.command_type = CommandType.EXECUTE_STATEMENT - should_retry, msg = server_directed_retry_policy.should_retry( - "POST", 429, has_retry_after=True - ) + policy = self._make_retry_policy(respect_server_retry_after_header=True) + policy._retry_start_time = time.time() + policy.command_type = CommandType.EXECUTE_STATEMENT + should_retry, msg = policy.should_retry("POST", 429, has_retry_after=True) assert should_retry is True def test_404_does_not_retry_for_any_command_type(self, retry_policy): From 945c14ab179ace306de8d9f6962d4c05e20bd3fe Mon Sep 17 00:00:00 2001 From: Shubham Dhal Date: Wed, 18 Mar 2026 17:47:33 +0530 Subject: [PATCH 4/6] Fix Black formatting for retry.py and utils.py Signed-off-by: Shubham Dhal --- src/databricks/sql/auth/retry.py | 5 ++++- src/databricks/sql/utils.py | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 454e2057c..140e7845b 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -390,7 +390,10 @@ def should_retry( # server explicitly signals it's safe via a Retry-After header. This prevents # duplicate side effects for non-idempotent operations. if self.respect_server_retry_after_header and not has_retry_after: - return (False, "respect_server_retry_after_header mode: no Retry-After header present") + return ( + False, + "respect_server_retry_after_header mode: no Retry-After header present", + ) # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry. if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS: diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 7311c30ba..93ab980f8 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -919,7 +919,9 @@ def build_client_context(server_hostname: str, version: str, **kwargs): ), retry_delay_default=kwargs.get("_retry_delay_default"), retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"), - respect_server_retry_after_header=kwargs.get("_respect_server_retry_after_header"), + respect_server_retry_after_header=kwargs.get( + "_respect_server_retry_after_header" + ), proxy_auth_method=kwargs.get("_proxy_auth_method"), pool_connections=kwargs.get("_pool_connections"), pool_maxsize=kwargs.get("_pool_maxsize"), From 616183f4bcb1abb7cfcc84de55aab7232f191f3e Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Fri, 20 Mar 2026 10:47:35 +0530 Subject: [PATCH 5/6] Apply CI workflow and test fixes from PR #745 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Update actions/checkout@v2 → v4 and actions/setup-python@v2 → v5 in run-unit-tests-with-arrow job - Pin Poetry version to 2.2.1 across all workflow files - Reduce long-running query test min duration from 3 to 2 minutes Co-authored-by: Isaac --- .github/workflows/code-quality-checks.yml | 8 ++++++-- .github/workflows/integration.yml | 1 + .github/workflows/publish-test.yml | 1 + .github/workflows/publish.yml | 1 + tests/e2e/common/large_queries_mixin.py | 6 +++--- 5 files changed, 12 insertions(+), 5 deletions(-) diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index 22db995c5..2d49707a5 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -35,6 +35,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true @@ -106,10 +107,10 @@ jobs: # check-out repo and set-up python #---------------------------------------------- - name: Check out repository - uses: actions/checkout@v2 + uses: actions/checkout@v4 - name: Set up python ${{ matrix.python-version }} id: setup-python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} #---------------------------------------------- @@ -118,6 +119,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true @@ -191,6 +193,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true @@ -243,6 +246,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 127c8ff4f..cf62d35d8 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -33,6 +33,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true diff --git a/.github/workflows/publish-test.yml b/.github/workflows/publish-test.yml index 2e6359a78..b7ffee9f4 100644 --- a/.github/workflows/publish-test.yml +++ b/.github/workflows/publish-test.yml @@ -21,6 +21,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index dde6cc2dc..c592756b8 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -23,6 +23,7 @@ jobs: - name: Install Poetry uses: snok/install-poetry@v1 with: + version: "2.2.1" virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true diff --git a/tests/e2e/common/large_queries_mixin.py b/tests/e2e/common/large_queries_mixin.py index dd7c56996..ad8538f8b 100644 --- a/tests/e2e/common/large_queries_mixin.py +++ b/tests/e2e/common/large_queries_mixin.py @@ -106,11 +106,11 @@ def test_query_with_large_narrow_result_set(self, extra_params): ], ) def test_long_running_query(self, extra_params): - """Incrementally increase query size until it takes at least 3 minutes, + """Incrementally increase query size until it takes at least 2 minutes, and asserts that the query completes successfully. """ minutes = 60 - min_duration = 3 * minutes + min_duration = 2 * minutes duration = -1 scale0 = 10000 @@ -136,5 +136,5 @@ def test_long_running_query(self, extra_params): duration = time.time() - start current_fraction = duration / min_duration print("Took {} s with scale factor={}".format(duration, scale_factor)) - # Extrapolate linearly to reach 3 min and add 50% padding to push over the limit + # Extrapolate linearly to reach 2 min and add 50% padding to push over the limit scale_factor = math.ceil(1.5 * scale_factor / current_fraction) From c96e14b17dc4966d045de8fe175fd2e81c504d0b Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Fri, 20 Mar 2026 11:00:44 +0530 Subject: [PATCH 6/6] Removed unnecessary tests --- tests/unit/test_retry.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index 7abab509d..635b611e2 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -167,14 +167,3 @@ def test_respect_server_retry_after__execute_statement_with_retry_after(self): should_retry, msg = policy.should_retry("POST", 429, has_retry_after=True) assert should_retry is True - def test_404_does_not_retry_for_any_command_type(self, retry_policy): - """Test that 404 never retries for any CommandType""" - retry_policy._retry_start_time = time.time() - - # Test for each CommandType - for command_type in CommandType: - retry_policy.command_type = command_type - should_retry, msg = retry_policy.should_retry("POST", 404) - - assert should_retry is False, f"404 should not retry for {command_type}" - assert "404" in msg or "NOT_FOUND" in msg