FEAT: add optional RetryPolicy for transient failures on connect() (GH-682)

CHANGELOG.md

@@ -55,6 +55,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
   for diagnostics (selected id, package, version, driver path, source, and
   whether it's frozen). This PR
   does not change the default provider or ship any Rust driver binaries.
+- **GH-682:** New optional `RetryPolicy` class and `retry_policy=` parameter on
+  `connect()` / `Connection(...)` that retries a connection attempt failing with
+  a transient SQLSTATE (login and connection timeouts, a lost link, `40001`,
+  `40003`) using exponential or fixed backoff, optional jitter and a delay cap.
+  `max_attempts` counts total tries including the first; without a policy
+  `connect()` behaves exactly as before.
 
 ### Changed
 - Connection strings and string connection parameters that contain a NUL

mssql_python/__init__.py

@@ -57,6 +57,9 @@
 # Token provider protocol (structural type for the token_provider= parameter)
 from .connection import TokenProvider
 
+# Retry policy for transient failures at connect() time (the retry_policy= parameter)
+from .retry import RetryPolicy
+
 # Connection String Handling
 from .connection_string_parser import _ConnectionStringParser
 from .connection_string_builder import _ConnectionStringBuilder
@@ -343,6 +346,8 @@ def _cleanup_connections():
     "TokenProvider",
     "Cursor",
     "Row",
+    # Retry policy
+    "RetryPolicy",
     # Settings
     "Settings",
     "get_settings",

mssql_python/connection.py

@@ -27,6 +27,7 @@
 from mssql_python.connection_string_parser import sanitize_connection_string
 from mssql_python.logging import logger
 from mssql_python import ddbc_bindings
+from mssql_python import retry
 from mssql_python.pooling import PoolingManager
 from mssql_python.odbc_provider import ProviderManager
 from mssql_python.exceptions import (
@@ -131,6 +132,26 @@ def _raise_connection_error(e: RuntimeError) -> None:
     ) from None
 
 
+def _sqlstate_from_runtime_error(e: RuntimeError) -> Optional[str]:
+    """Return the SQLSTATE carried by a RuntimeError from the C++ pybind layer.
+
+    Connection::checkError() throws "SQLSTATE:XXXXX:<odbc_message>". Only a code of exactly five
+    characters is returned; a message without the prefix, or with an empty or truncated code,
+    yields None so the caller treats the failure as not retriable.
+
+    Args:
+        e (RuntimeError): The exception raised by the native connection.
+
+    Returns:
+        Optional[str]: The SQLSTATE, or None.
+    """
+    match = _SQLSTATE_RE.match(str(e))
+    if match is None:
+        return None
+    sqlstate = match.group(1)
+    return sqlstate if len(sqlstate) == 5 else None
+
+
 def _validate_utf16_wchar_compatibility(
     encoding: str, wchar_type: int, context: str = "SQL_WCHAR"
 ) -> None:
@@ -278,6 +299,7 @@ def __init__(
         timeout: int = 0,
         native_uuid: Optional[bool] = None,
         token_provider: Optional["TokenProvider"] = None,
+        retry_policy: Optional[retry.RetryPolicy] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -344,6 +366,19 @@ def __init__(
                     Interactive credentials (e.g. ``InteractiveBrowserCredential``) block
                     ``connect()`` until the user completes sign-in; prefer non-interactive
                     credentials in server contexts.
+            retry_policy (RetryPolicy, optional): Policy for retrying the native connect when
+                it fails with a transient SQLSTATE (a login or connection timeout, a lost link
+                and similar; see ``mssql_python.retry.DEFAULT_RETRIABLE_SQLSTATES``). None
+                (default) makes a single attempt, exactly as before. The connection string is
+                parsed once, before the first attempt, and a token acquired on the Python side
+                (``token_provider=``, ``Authentication=ActiveDirectoryDefault`` or a raw
+                ``attrs_before`` token) is acquired once and reused by every attempt. For
+                managed identity, interactive and device code authentication the native layer
+                asks the deferred token factory for a token on each physical connect, so a
+                retried attempt may acquire a fresh one. The login timeout bounds each attempt
+                separately, so the total wall clock time is roughly the attempt timeouts plus
+                the delays. Each retry is logged at warning level through the driver logger,
+                which shows it once ``setup_logging()`` has been called.
             **kwargs: Additional key/value pairs for the connection string.
 
         Returns:
@@ -354,6 +389,7 @@ def __init__(
                 source, or lacking a valid ``.get_token`` method), or the credential returns
                 no valid token.
             OperationalError: If acquiring a token from ``token_provider`` fails.
+            TypeError: If ``retry_policy`` is neither None nor a ``RetryPolicy``.
             ValueError: If the connection string is invalid or connection fails.
 
         This method sets up the initial state for the connection object,
@@ -375,6 +411,16 @@ def __init__(
             raise ValueError("native_uuid must be a boolean value or None")
         self._native_uuid = native_uuid
 
+        # Check the retry policy type up front, before the connection string is parsed or a
+        # token is acquired, so a wrong value fails fast with no network work. It is kept on
+        # the connection so cursor level retries can later pick it up as their default.
+        if retry_policy is not None and not isinstance(retry_policy, retry.RetryPolicy):
+            raise TypeError(
+                "retry_policy must be a RetryPolicy instance or None, "
+                f"got {type(retry_policy).__name__}"
+            )
+        self._retry_policy: Optional[retry.RetryPolicy] = retry_policy
+
         self.connection_str, parsed_params = self._construct_connection_string(
             connection_str, **kwargs
         )
@@ -741,16 +787,42 @@ def _token_factory():
         _provider = ProviderManager.ensure_available()
         ddbc_bindings._set_odbc_provider(_provider)
 
-        try:
-            self._conn = ddbc_bindings.Connection(
-                self.connection_str,
-                self._pooling,
-                self._attrs_before,
-                self._pool_key,
-                self._token_factory,
-            )
-        except RuntimeError as e:
-            _raise_connection_error(e)
+        # A retry policy wraps only the native connect. Everything above (connection string
+        # parsing, the attrs_before copy, any token acquired on the Python side) has already
+        # happened once, so every attempt reuses the same inputs. A deferred token factory is
+        # still invoked by the native layer on each physical connect, so those paths may acquire
+        # a fresh token per attempt. Without a policy this is a single attempt, exactly the
+        # behaviour before retry_policy existed.
+        max_attempts = retry_policy.max_attempts if retry_policy is not None else 1
+        for attempt in range(1, max_attempts + 1):
+            try:
+                self._conn = ddbc_bindings.Connection(
+                    self.connection_str,
+                    self._pooling,
+                    self._attrs_before,
+                    self._pool_key,
+                    self._token_factory,
+                )
+                break
+            except RuntimeError as e:
+                sqlstate = _sqlstate_from_runtime_error(e)
+                if (
+                    retry_policy is not None
+                    and attempt < max_attempts
+                    and retry_policy.is_retriable(sqlstate)
+                ):
+                    delay = retry_policy.compute_delay(attempt)
+                    logger.warning(
+                        "Connection attempt %d of %d failed with SQLSTATE %s; "
+                        "retry in %.2f seconds",
+                        attempt,
+                        max_attempts,
+                        sqlstate,
+                        delay,
+                    )
+                    retry._sleep(delay)  # pylint: disable=protected-access
+                    continue
+                _raise_connection_error(e)
         self.setautocommit(autocommit)
 
         # Register this connection for cleanup before Python shutdown

mssql_python/db_connection.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, Optional, Union
 
 from mssql_python.connection import Connection, TokenProvider
+from mssql_python.retry import RetryPolicy
 
 
 def connect(
@@ -16,6 +17,7 @@ def connect(
     timeout: int = 0,
     native_uuid: Optional[bool] = None,
     token_provider: Optional[TokenProvider] = None,
+    retry_policy: Optional[RetryPolicy] = None,
     **kwargs: Any,
 ) -> Connection:
     """
@@ -69,6 +71,16 @@ def connect(
             (``https://database.windows.net/.default``). Sovereign clouds (Azure US
             Government, Azure China, Azure Germany) are **out of scope** — acquire the token
             yourself and pass it via ``attrs_before[SQL_COPT_SS_ACCESS_TOKEN]`` instead.
+        retry_policy (RetryPolicy, optional): Policy for retrying the connection attempt when
+            it fails with a transient SQLSTATE such as a login timeout or a lost link. None
+            (default) makes a single attempt, exactly as before. See ``RetryPolicy`` for the
+            settings and ``mssql_python.retry.DEFAULT_RETRIABLE_SQLSTATES`` for the codes
+            retried by default.
+
+            Example::
+
+                policy = mssql_python.RetryPolicy(max_attempts=5, base_delay=0.5)
+                conn = mssql_python.connect("Server=s;Database=d", retry_policy=policy)
     Keyword Args:
         **kwargs: Additional key/value pairs for the connection string.
     Below attributes are not implemented in the internal driver:
@@ -81,6 +93,7 @@ def connect(
     Raises:
         DatabaseError: If there is an error while trying to connect to the database.
         InterfaceError: If there is an error related to the database interface.
+        TypeError: If ``retry_policy`` is neither None nor a ``RetryPolicy``.
 
     This function provides a way to create a new connection object, which can then
     be used to perform database operations such as executing queries, committing
@@ -93,6 +106,7 @@ def connect(
         timeout=timeout,
         native_uuid=native_uuid,
         token_provider=token_provider,
+        retry_policy=retry_policy,
         **kwargs,
     )
     return conn

mssql_python/mssql_python.pyi

@@ -7,6 +7,7 @@ Type stubs for mssql_python package - based on actual public API
 from typing import (
     Any,
     Dict,
+    FrozenSet,
     List,
     Mapping,
     Optional,
@@ -273,6 +274,40 @@ class _ArrowReader:
         use_internal_transaction: bool = False,
     ) -> Dict[str, Any]: ...
 
+# Retry Policy for transient failures at connect() time
+class RetryPolicy:
+    """
+    Describes how connect() retries a connection attempt that fails with a transient error.
+
+    Pass an instance as the retry_policy= argument of connect() or Connection().
+    max_attempts is the total number of tries including the first; 1 means never retry.
+    """
+
+    @property
+    def max_attempts(self) -> int: ...
+    @property
+    def backoff(self) -> str: ...
+    @property
+    def base_delay(self) -> float: ...
+    @property
+    def max_delay(self) -> float: ...
+    @property
+    def jitter(self) -> bool: ...
+    @property
+    def retriable_sqlstates(self) -> FrozenSet[str]: ...
+    def __init__(
+        self,
+        max_attempts: int = 3,
+        backoff: str = "exponential",
+        base_delay: float = 1.0,
+        max_delay: float = 30.0,
+        jitter: bool = True,
+        retriable_sqlstates: Optional[Iterable[str]] = None,
+    ) -> None: ...
+    def is_retriable(self, sqlstate: Optional[str]) -> bool: ...
+    def compute_delay(self, attempt: int) -> float: ...
+    def __repr__(self) -> str: ...
+
 # DB-API 2.0 Connection Object
 # https://www.python.org/dev/peps/pep-0249/#connection-objects
 class Connection:
@@ -312,6 +347,7 @@ class Connection:
         attrs_before: Optional[Dict[int, Union[int, str, bytes]]] = None,
         timeout: int = 0,
         native_uuid: Optional[bool] = None,
+        retry_policy: Optional[RetryPolicy] = None,
         **kwargs: Any,
     ) -> None: ...
 
@@ -357,6 +393,7 @@ def connect(
     attrs_before: Optional[Dict[int, Union[int, str, bytes]]] = None,
     timeout: int = 0,
     native_uuid: Optional[bool] = None,
+    retry_policy: Optional[RetryPolicy] = None,
     **kwargs: Any,
 ) -> Connection: ...