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 <shubham.dhal@databricks.com>

Author: sd-db

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,
@@ -80,6 +81,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

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:
@@ -388,6 +393,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, was an ExecuteStatement and the command may have reached the server
         if (
             self.command_type == CommandType.EXECUTE_STATEMENT
@@ -430,7 +441,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)

src/databricks/sql/backend/sea/utils/http_client.py

@@ -92,6 +92,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)
@@ -116,6 +119,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:

src/databricks/sql/backend/thrift_backend.py

@@ -191,6 +191,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 = {}
 
@@ -217,6 +220,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,
             )
 

src/databricks/sql/common/unified_http_client.py

@@ -135,6 +135,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

src/databricks/sql/utils.py

@@ -977,6 +977,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"),

tests/unit/test_retry.py

@@ -84,6 +84,114 @@ 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
+    ):
+        """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()

tests/unit/test_unified_http_client.py

@@ -37,6 +37,7 @@ def client_context(self):
         context.retry_stop_after_attempts_duration = 300.0
         context.retry_delay_default = 5.0
         context.retry_dangerous_codes = []
+        context.retry_server_directed_only = False
         context.proxy_auth_method = None
         context.pool_connections = 10
         context.pool_maxsize = 20
@@ -48,16 +49,19 @@ def http_client(self, client_context):
         """Create UnifiedHttpClient instance."""
         return UnifiedHttpClient(client_context)
 
-    @pytest.mark.parametrize("status_code,path", [
-        (429, "reason.response"),
-        (503, "reason.response"),
-        (500, "direct_response"),
-    ])
+    @pytest.mark.parametrize(
+        "status_code,path",
+        [
+            (429, "reason.response"),
+            (503, "reason.response"),
+            (500, "direct_response"),
+        ],
+    )
     def test_max_retry_error_with_status_codes(self, http_client, status_code, path):
         """Test MaxRetryError with various status codes and response paths."""
         mock_pool = Mock()
         max_retry_error = MaxRetryError(pool=mock_pool, url="http://test.com")
-        
+
         if path == "reason.response":
             max_retry_error.reason = Mock()
             max_retry_error.reason.response = Mock()
@@ -79,12 +83,21 @@ def test_max_retry_error_with_status_codes(self, http_client, status_code, path)
             assert "http-code" in error.context
             assert error.context["http-code"] == status_code
 
-    @pytest.mark.parametrize("setup_func", [
-        lambda e: None,  # No setup - error with no attributes
-        lambda e: setattr(e, "reason", None),  # reason=None
-        lambda e: (setattr(e, "reason", Mock()), setattr(e.reason, "response", None)),  # reason.response=None
-        lambda e: (setattr(e, "reason", Mock()), setattr(e.reason, "response", Mock(spec=[]))),  # No status attr
-    ])
+    @pytest.mark.parametrize(
+        "setup_func",
+        [
+            lambda e: None,  # No setup - error with no attributes
+            lambda e: setattr(e, "reason", None),  # reason=None
+            lambda e: (
+                setattr(e, "reason", Mock()),
+                setattr(e.reason, "response", None),
+            ),  # reason.response=None
+            lambda e: (
+                setattr(e, "reason", Mock()),
+                setattr(e.reason, "response", Mock(spec=[])),
+            ),  # No status attr
+        ],
+    )
     def test_max_retry_error_missing_status(self, http_client, setup_func):
         """Test MaxRetryError without status code (no crash, empty context)."""
         mock_pool = Mock()
@@ -104,12 +117,12 @@ def test_max_retry_error_prefers_reason_response(self, http_client):
         """Test that e.reason.response.status is preferred over e.response.status."""
         mock_pool = Mock()
         max_retry_error = MaxRetryError(pool=mock_pool, url="http://test.com")
-        
+
         # Set both structures with different status codes
         max_retry_error.reason = Mock()
         max_retry_error.reason.response = Mock()
         max_retry_error.reason.response.status = 429  # Should use this
-        
+
         max_retry_error.response = Mock()
         max_retry_error.response.status = 500  # Should be ignored