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

Author: sd-db

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

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:

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:

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,
             )
 

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

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"),

tests/unit/test_retry.py

@@ -7,16 +7,21 @@
 
 
 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,
             stop_after_attempts_duration=900,
             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,112 +89,82 @@ 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(
             "POST", 429, has_retry_after=False
         )
         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):