diff --git a/CHANGELOG.md b/CHANGELOG.md index f4ebfb81..97c43baf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/mssql_python/__init__.py b/mssql_python/__init__.py index b5a4fe84..7eea5468 100644 --- a/mssql_python/__init__.py +++ b/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", diff --git a/mssql_python/connection.py b/mssql_python/connection.py index 6ed9f1f9..4b92cda9 100644 --- a/mssql_python/connection.py +++ b/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:". 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 diff --git a/mssql_python/db_connection.py b/mssql_python/db_connection.py index ec706709..73772cff 100644 --- a/mssql_python/db_connection.py +++ b/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 diff --git a/mssql_python/mssql_python.pyi b/mssql_python/mssql_python.pyi index c8cc076d..39cc0fe7 100644 --- a/mssql_python/mssql_python.pyi +++ b/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: ... diff --git a/mssql_python/retry.py b/mssql_python/retry.py new file mode 100644 index 00000000..94610b64 --- /dev/null +++ b/mssql_python/retry.py @@ -0,0 +1,220 @@ +""" +Copyright (c) Microsoft Corporation. +Licensed under the MIT license. +This module defines the RetryPolicy class, which describes how connect() retries a connection +attempt that fails with a transient error. +""" + +import math +import random +import time +from typing import FrozenSet, Iterable, Optional + +# Seams for tests. Both are looked up on this module at call time, so a test can replace them +# and assert on the exact delay sequence without sleeping or depending on the random source. +_sleep = time.sleep +_random = random.random + +# SQLSTATEs the driver treats as transient at connect time. These are the seven transient codes +# from the retry logic page for the driver on Microsoft Learn +# (https://learn.microsoft.com/sql/connect/python/mssql-python/retry-logic), applied here to +# the connect attempt: HYT00 and HYT01 (a timeout), 08001, 08S01 and 08007 (the link could not +# be established or was lost), 40001 (serialization failure) and 40003 (statement completion +# unknown). 08004, "Server rejected the connection", is deliberately excluded: the server +# answered and refused, so the same request is not going to be accepted on the next try. +DEFAULT_RETRIABLE_SQLSTATES: FrozenSet[str] = frozenset( + {"HYT00", "HYT01", "08001", "08S01", "08007", "40001", "40003"} +) + +_BACKOFF_STRATEGIES = ("exponential", "fixed") +_SQLSTATE_LENGTH = 5 + + +def _is_finite_number(value: object) -> bool: + """Return True for a finite int or float that is not a bool.""" + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) + + +def _normalize_sqlstates(codes: Optional[Iterable[str]]) -> FrozenSet[str]: + """Validate and upper case a caller supplied set of SQLSTATE codes. + + Args: + codes (iterable of str, optional): SQLSTATE codes, or None for the driver default set. + + Returns: + FrozenSet[str]: The upper cased codes, or ``DEFAULT_RETRIABLE_SQLSTATES`` for None. + + Raises: + ValueError: If ``codes`` is a single string, or any code is not a string of exactly + five characters. + """ + if codes is None: + return DEFAULT_RETRIABLE_SQLSTATES + if isinstance(codes, (str, bytes)): + raise ValueError( + "retriable_sqlstates must be an iterable of SQLSTATE strings, not a single string" + ) + normalized = set() + for code in codes: + if not isinstance(code, str) or len(code) != _SQLSTATE_LENGTH: + raise ValueError( + f"each SQLSTATE must be a string of exactly {_SQLSTATE_LENGTH} characters, " + f"got {code!r}" + ) + normalized.add(code.upper()) + return frozenset(normalized) + + +class RetryPolicy: + """Describes how ``connect()`` retries a connection attempt that fails with a transient error. + + A policy is optional: ``connect()`` and ``Connection()`` make a single attempt unless one is + passed as ``retry_policy=``. When the native connect raises with a SQLSTATE in + ``retriable_sqlstates``, the driver waits for ``compute_delay(attempt)`` seconds and tries + again, up to ``max_attempts`` tries in total. Any other failure is raised at once, as the + same exception type it has always been. + + Every setting is validated once in ``__init__`` and exposed through a property with no + setter, so an instance cannot be changed after construction and the same policy can be + shared by any number of connections. + + Attributes: + max_attempts (int): Total number of tries, including the first. 1 means never retry. + backoff (str): "exponential" doubles the delay after each failed attempt, "fixed" + waits ``base_delay`` every time. + base_delay (float): Delay in seconds before the second attempt. + max_delay (float): Upper bound in seconds for any single delay, jitter included. + jitter (bool): When True each delay is scaled by a factor drawn uniformly from + [0.5, 1.5) so that many clients do not reconnect in lockstep. + retriable_sqlstates (frozenset): The SQLSTATE codes that are retried, uppercased and + each exactly five characters. Defaults to ``DEFAULT_RETRIABLE_SQLSTATES``; a custom + set replaces the default entirely rather than extending it. + + Example: + >>> import mssql_python as ms + >>> policy = ms.RetryPolicy(max_attempts=5, base_delay=0.5, max_delay=10.0) + >>> conn = ms.connect("Server=myserver;Database=mydb", retry_policy=policy) + """ + + 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: + """Validate the settings and build the policy. + + Args: + max_attempts (int): Total number of tries including the first; at least 1. + backoff (str): "exponential" or "fixed". + base_delay (float): Seconds to wait before the second attempt; zero or more. + max_delay (float): Cap in seconds for every delay; at least ``base_delay``. + jitter (bool): Scale each delay by a random factor in [0.5, 1.5). + retriable_sqlstates (iterable of str, optional): SQLSTATE codes to retry. None + selects ``DEFAULT_RETRIABLE_SQLSTATES``. Codes are upper cased. + + Raises: + ValueError: If any setting is out of range or of the wrong type. + """ + if isinstance(max_attempts, bool) or not isinstance(max_attempts, int): + raise ValueError("max_attempts must be an integer of at least 1") + if max_attempts < 1: + raise ValueError("max_attempts must be an integer of at least 1") + if backoff not in _BACKOFF_STRATEGIES: + raise ValueError("backoff must be one of 'exponential' or 'fixed'") + if not _is_finite_number(base_delay) or base_delay < 0: + raise ValueError("base_delay must be a finite number of zero or more seconds") + if not _is_finite_number(max_delay) or max_delay < base_delay: + raise ValueError("max_delay must be a finite number of at least base_delay seconds") + if not isinstance(jitter, bool): + raise ValueError("jitter must be True or False") + + self._max_attempts: int = max_attempts + self._backoff: str = backoff + self._base_delay: float = float(base_delay) + self._max_delay: float = float(max_delay) + self._jitter: bool = jitter + self._retriable_sqlstates: FrozenSet[str] = _normalize_sqlstates(retriable_sqlstates) + + @property + def max_attempts(self) -> int: + """Total number of tries, including the first.""" + return self._max_attempts + + @property + def backoff(self) -> str: + """Backoff strategy, "exponential" or "fixed".""" + return self._backoff + + @property + def base_delay(self) -> float: + """Delay in seconds before the second attempt.""" + return self._base_delay + + @property + def max_delay(self) -> float: + """Upper bound in seconds for any single delay, jitter included.""" + return self._max_delay + + @property + def jitter(self) -> bool: + """Whether each delay is scaled by a random factor in [0.5, 1.5).""" + return self._jitter + + @property + def retriable_sqlstates(self) -> FrozenSet[str]: + """The SQLSTATE codes this policy retries.""" + return self._retriable_sqlstates + + def is_retriable(self, sqlstate: Optional[str]) -> bool: + """Return True when ``sqlstate`` is one of the codes this policy retries. + + Args: + sqlstate (str, optional): SQLSTATE code from the failed attempt, or None when the + failure carried no SQLSTATE. None is never retriable. + + Returns: + bool: True only when the upper cased code is in ``retriable_sqlstates``. + """ + if not isinstance(sqlstate, str): + return False + return sqlstate.upper() in self.retriable_sqlstates + + def compute_delay(self, attempt: int) -> float: + """Return how long to wait, in seconds, after a failed attempt. + + Args: + attempt (int): Index, counting from 1, of the attempt that just failed, so the + delay before the second attempt is ``compute_delay(1)``. + + Returns: + float: Seconds to wait, never negative and never above ``max_delay``. + + Raises: + ValueError: If ``attempt`` is less than 1. + """ + if isinstance(attempt, bool) or not isinstance(attempt, int) or attempt < 1: + raise ValueError("attempt must be an integer of at least 1") + delay = self.base_delay + if self.backoff == "exponential": + # Double once per failed attempt and stop as soon as the cap is reached, so a large + # attempt number can never overflow the way a direct power of two would. + doublings = attempt - 1 + while doublings > 0 and 0.0 < delay < self.max_delay: + delay *= 2.0 + doublings -= 1 + delay = min(delay, self.max_delay) + if self.jitter: + delay = min(delay * (0.5 + _random()), self.max_delay) + return delay + + def __repr__(self) -> str: + """Return a constructor style representation of the policy.""" + return ( + f"RetryPolicy(max_attempts={self.max_attempts!r}, backoff={self.backoff!r}, " + f"base_delay={self.base_delay!r}, max_delay={self.max_delay!r}, " + f"jitter={self.jitter!r}, retriable_sqlstates={sorted(self.retriable_sqlstates)!r})" + ) diff --git a/tests/test_027_retry_policy.py b/tests/test_027_retry_policy.py new file mode 100644 index 00000000..8f760b0f --- /dev/null +++ b/tests/test_027_retry_policy.py @@ -0,0 +1,406 @@ +""" +Tests for the optional retry policy on connect(), added for +https://github.com/microsoft/mssql-python/issues/682. + +No test here needs a server. The native connection constructor is replaced with a fake that +fails a chosen number of times, and the retry module's sleep and random seams are replaced so +nothing sleeps and every delay sequence is asserted exactly. Neither the db_connection nor the +cursor fixture is requested, so the file runs with DB_CONNECTION_STRING unset. +""" + +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +import mssql_python +import mssql_python.connection +import mssql_python.logging +import mssql_python.retry +from mssql_python import Connection, RetryPolicy, connect +from mssql_python.exceptions import OperationalError, ProgrammingError +from mssql_python.retry import DEFAULT_RETRIABLE_SQLSTATES + +CONN_STR = "Server=testserver;Database=mydb;Trusted_Connection=yes;" +DRIVER_PREFIX = "[Microsoft][ODBC Driver 18 for SQL Server]" +LINK_FAILURE = "SQLSTATE:08S01:" + DRIVER_PREFIX + "Communication link failure" +LOGIN_FAILURE = "SQLSTATE:28000:" + DRIVER_PREFIX + "Login failed for user 'baduser'." +THE_SEVEN = ("HYT00", "HYT01", "08001", "08S01", "08007", "40001", "40003") + + +class FakeNativeConnection: + """Stand in for ddbc_bindings.Connection that fails a set number of times, then succeeds. + + Every call records its positional arguments, so a test can assert how many attempts were + made and that each attempt received exactly the same inputs. + """ + + def __init__(self, failures=0, message=LINK_FAILURE): + self.failures = failures + self.message = message + self.calls = [] + + def __call__(self, *args, **kwargs): + # Snapshot any dict argument, so a mutation between attempts shows up as a difference + # between recorded calls instead of the same object being compared with itself. + self.calls.append(tuple(dict(arg) if isinstance(arg, dict) else arg for arg in args)) + if len(self.calls) <= self.failures: + raise RuntimeError(self.message) + native = MagicMock() + native.get_autocommit.return_value = False + return native + + +class CountingTokenProvider: + """Minimal token_provider whose get_token() counts how often a token is requested.""" + + def __init__(self): + self.calls = 0 + + def get_token(self, scope): + self.calls += 1 + return SimpleNamespace(token="header.payload.signature", expires_on=None) + + +class RecordingHandler(logging.Handler): + """Collects the formatted messages the driver logger emits.""" + + def __init__(self): + super().__init__() + self.messages = [] + + def emit(self, record): + self.messages.append((record.levelno, record.getMessage())) + + +@pytest.fixture(autouse=True) +def sleeps(monkeypatch): + """Replace the retry module's sleep with a recorder so no test ever waits.""" + recorded = [] + monkeypatch.setattr(mssql_python.retry, "_sleep", recorded.append) + return recorded + + +@pytest.fixture +def native(monkeypatch): + """Install a FakeNativeConnection in place of the pybind constructor.""" + fake = FakeNativeConnection() + monkeypatch.setattr(mssql_python.connection.ddbc_bindings, "Connection", fake) + return fake + + +@pytest.fixture +def driver_log(): + """Attach a recording handler to the driver logger for the duration of a test. + + The underlying stdlib logger sits at CRITICAL until setup_logging() is called, so its level + is lowered to WARNING here and restored afterwards; nothing else about logging is changed. + """ + stdlib_logger = logging.getLogger("mssql_python") + previous_level = stdlib_logger.level + stdlib_logger.setLevel(logging.WARNING) + handler = RecordingHandler() + mssql_python.logging.logger.addHandler(handler) + try: + yield handler + finally: + mssql_python.logging.logger.removeHandler(handler) + stdlib_logger.setLevel(previous_level) + + +def test_no_policy_makes_a_single_attempt_and_raises_as_before(native, sleeps): + native.failures = 1 + with pytest.raises(OperationalError) as exc_info: + connect(CONN_STR) + assert len(native.calls) == 1 + assert sleeps == [] + assert exc_info.value.driver_error == "Communication link failure" + assert "Communication link failure" in exc_info.value.ddbc_error + assert not isinstance(exc_info.value, RuntimeError) + + +def test_no_policy_is_stored_as_none(native): + conn = connect(CONN_STR) + assert conn._retry_policy is None + assert len(native.calls) == 1 + + +def test_policy_retries_transient_failure_until_success(native, sleeps): + native.failures = 2 + policy = RetryPolicy(max_attempts=3, jitter=False) + conn = connect(CONN_STR, retry_policy=policy) + assert len(native.calls) == 3 + assert sleeps == [1.0, 2.0] + assert conn._retry_policy is policy + # Every attempt is made with the same connection string, attributes, pool key and factory. + assert all(call == native.calls[0] for call in native.calls) + + +def test_policy_does_not_retry_permanent_failure(native, sleeps): + native.failures = 1 + native.message = LOGIN_FAILURE + with pytest.raises(OperationalError) as exc_info: + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert len(native.calls) == 1 + assert sleeps == [] + assert exc_info.value.driver_error == "Invalid authorization specification" + + +def test_policy_exhausts_attempts_and_raises_the_mapped_type(native, sleeps): + native.failures = 3 + with pytest.raises(OperationalError) as exc_info: + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert len(native.calls) == 3 + assert sleeps == [1.0, 2.0] + assert type(exc_info.value) is OperationalError + assert exc_info.value.driver_error == "Communication link failure" + assert "Communication link failure" in exc_info.value.ddbc_error + assert not isinstance(exc_info.value, RuntimeError) + + +@pytest.mark.parametrize( + "message", + [ + pytest.param(DRIVER_PREFIX + "Connection handle not allocated", id="no_prefix"), + pytest.param("SQLSTATE::" + DRIVER_PREFIX + "Invalid handle!", id="empty_code"), + ], +) +def test_policy_does_not_retry_error_without_a_sqlstate(native, sleeps, message): + native.failures = 1 + native.message = message + with pytest.raises(OperationalError) as exc_info: + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert len(native.calls) == 1 + assert sleeps == [] + assert exc_info.value.driver_error == "Connection operation failed" + + +def test_default_set_is_exactly_the_seven_transient_codes(): + assert DEFAULT_RETRIABLE_SQLSTATES == frozenset(THE_SEVEN) + assert RetryPolicy().retriable_sqlstates is DEFAULT_RETRIABLE_SQLSTATES + + +@pytest.mark.parametrize("sqlstate", THE_SEVEN) +def test_default_policy_retries_each_transient_sqlstate(native, sleeps, sqlstate): + assert RetryPolicy().is_retriable(sqlstate) + assert RetryPolicy().is_retriable(sqlstate.lower()) + native.failures = 1 + native.message = "SQLSTATE:" + sqlstate + ":" + DRIVER_PREFIX + "transient failure" + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=2, jitter=False)) + assert len(native.calls) == 2 + assert sleeps == [1.0] + + +@pytest.mark.parametrize( + "sqlstate, expected", + [ + ("08004", OperationalError), + ("28000", OperationalError), + ("42000", ProgrammingError), + ], +) +def test_default_policy_does_not_retry_permanent_sqlstate(native, sleeps, sqlstate, expected): + assert not RetryPolicy().is_retriable(sqlstate) + native.failures = 1 + native.message = "SQLSTATE:" + sqlstate + ":" + DRIVER_PREFIX + "permanent failure" + with pytest.raises(expected): + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert len(native.calls) == 1 + assert sleeps == [] + + +@pytest.mark.parametrize("sqlstate", [None, "", "08S0", "08S011"]) +def test_is_retriable_rejects_missing_or_malformed_codes(sqlstate): + assert not RetryPolicy().is_retriable(sqlstate) + + +def test_custom_sqlstates_replace_the_default_set(native, sleeps): + policy = RetryPolicy(max_attempts=2, jitter=False, retriable_sqlstates={"28000"}) + assert policy.retriable_sqlstates == frozenset({"28000"}) + assert policy.is_retriable("28000") + assert not policy.is_retriable("08S01") + native.failures = 1 + native.message = LOGIN_FAILURE + connect(CONN_STR, retry_policy=policy) + assert len(native.calls) == 2 + assert sleeps == [1.0] + + +def test_custom_sqlstates_do_not_retry_a_default_code(native, sleeps): + policy = RetryPolicy(max_attempts=2, jitter=False, retriable_sqlstates={"28000"}) + native.failures = 1 + with pytest.raises(OperationalError) as exc_info: + connect(CONN_STR, retry_policy=policy) + assert len(native.calls) == 1 + assert sleeps == [] + assert exc_info.value.driver_error == "Communication link failure" + + +def test_custom_sqlstates_are_upper_cased_and_accept_any_iterable(): + policy = RetryPolicy(retriable_sqlstates=["08s01", "hyt00"]) + assert policy.retriable_sqlstates == frozenset({"08S01", "HYT00"}) + assert RetryPolicy(retriable_sqlstates=()).retriable_sqlstates == frozenset() + + +def test_exponential_delay_doubles_and_is_capped(): + policy = RetryPolicy(max_attempts=6, base_delay=1.0, max_delay=5.0, jitter=False) + assert [policy.compute_delay(n) for n in range(1, 6)] == [1.0, 2.0, 4.0, 5.0, 5.0] + + +def test_exponential_delay_with_a_huge_attempt_number_stays_at_the_cap(): + policy = RetryPolicy(jitter=False) + assert policy.compute_delay(5000) == 30.0 + assert RetryPolicy(base_delay=0.0, jitter=False).compute_delay(5000) == 0.0 + + +def test_fixed_delay_is_constant(): + policy = RetryPolicy(backoff="fixed", base_delay=0.25, max_delay=5.0, jitter=False) + assert [policy.compute_delay(n) for n in range(1, 5)] == [0.25, 0.25, 0.25, 0.25] + + +def test_jitter_scales_the_delay_and_never_exceeds_the_cap(monkeypatch): + policy = RetryPolicy(base_delay=1.0, max_delay=5.0, jitter=True) + monkeypatch.setattr(mssql_python.retry, "_random", lambda: 0.0) + assert [policy.compute_delay(n) for n in (1, 2, 3)] == [0.5, 1.0, 2.0] + monkeypatch.setattr(mssql_python.retry, "_random", lambda: 1.0) + assert [policy.compute_delay(n) for n in (1, 2, 3, 4)] == [1.5, 3.0, 5.0, 5.0] + + +def test_jittered_delays_are_used_when_retrying(native, sleeps, monkeypatch): + monkeypatch.setattr(mssql_python.retry, "_random", lambda: 0.0) + native.failures = 2 + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3)) + assert sleeps == [0.5, 1.0] + + +@pytest.mark.parametrize("attempt", [0, -1, 1.0, True]) +def test_compute_delay_rejects_an_invalid_attempt_number(attempt): + with pytest.raises(ValueError): + RetryPolicy().compute_delay(attempt) + + +def test_default_settings_match_the_issue_proposal(): + policy = RetryPolicy() + assert policy.max_attempts == 3 + assert policy.backoff == "exponential" + assert policy.base_delay == 1.0 + assert policy.max_delay == 30.0 + assert policy.jitter is True + assert policy.retriable_sqlstates == DEFAULT_RETRIABLE_SQLSTATES + assert repr(policy).startswith("RetryPolicy(max_attempts=3, backoff='exponential'") + assert "08S01" in repr(policy) + + +@pytest.mark.parametrize( + "name, value", + [ + ("max_attempts", 0), + ("backoff", "fixed"), + ("base_delay", 2.0), + ("max_delay", 60.0), + ("jitter", False), + ("retriable_sqlstates", frozenset({"28000"})), + ], +) +def test_policy_settings_cannot_be_changed_after_construction(name, value): + policy = RetryPolicy() + with pytest.raises(AttributeError): + setattr(policy, name, value) + assert getattr(policy, name) == getattr(RetryPolicy(), name) + + +def test_single_attempt_policy_never_retries(native, sleeps): + native.failures = 1 + with pytest.raises(OperationalError): + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=1)) + assert len(native.calls) == 1 + assert sleeps == [] + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"max_attempts": 0}, id="max_attempts_zero"), + pytest.param({"max_attempts": True}, id="max_attempts_bool"), + pytest.param({"max_attempts": 2.0}, id="max_attempts_float"), + pytest.param({"backoff": "linear"}, id="backoff_linear"), + pytest.param({"base_delay": -1.0}, id="base_delay_negative"), + pytest.param({"base_delay": float("nan")}, id="base_delay_nan"), + pytest.param({"base_delay": 2.0, "max_delay": 1.0}, id="max_delay_below_base"), + pytest.param({"max_delay": float("inf")}, id="max_delay_infinite"), + pytest.param({"jitter": 1}, id="jitter_not_bool"), + pytest.param({"retriable_sqlstates": ["08S0"]}, id="sqlstate_four_chars"), + pytest.param({"retriable_sqlstates": "08S01"}, id="sqlstate_bare_string"), + pytest.param({"retriable_sqlstates": [8001]}, id="sqlstate_not_a_string"), + ], +) +def test_invalid_settings_raise_value_error(kwargs): + with pytest.raises(ValueError): + RetryPolicy(**kwargs) + + +def test_connect_rejects_a_value_that_is_not_a_policy(native, sleeps): + with pytest.raises(TypeError): + connect(CONN_STR, retry_policy="nope") + with pytest.raises(TypeError): + Connection(CONN_STR, retry_policy={"max_attempts": 3}) + assert native.calls == [] + assert sleeps == [] + + +def test_connect_passes_the_policy_through_to_the_connection(native): + policy = RetryPolicy(max_attempts=2) + conn = connect(CONN_STR, retry_policy=policy) + assert conn._retry_policy is policy + assert len(native.calls) == 1 + + +def test_token_is_acquired_once_across_attempts(native, sleeps): + native.failures = 2 + provider = CountingTokenProvider() + connect( + "Server=testserver;Database=mydb;", + token_provider=provider, + retry_policy=RetryPolicy(max_attempts=3, jitter=False), + ) + assert len(native.calls) == 3 + assert provider.calls == 1 + assert sleeps == [1.0, 2.0] + + +def test_retry_log_lines_name_the_attempt_and_omit_the_connection_string( + native, sleeps, driver_log +): + native.failures = 3 + with pytest.raises(OperationalError): + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + warnings = [msg for level, msg in driver_log.messages if level == logging.WARNING] + errors = [msg for level, msg in driver_log.messages if level == logging.ERROR] + assert len(warnings) == 2 + assert "attempt 1 of 3" in warnings[0] and "08S01" in warnings[0] + assert "attempt 2 of 3" in warnings[1] and "2.00 seconds" in warnings[1] + # The final failure logs only the one error line _raise_connection_error has always written. + assert len(errors) == 1 + assert "Connection attempt" not in errors[0] + retry_lines = [msg for _, msg in driver_log.messages if "Connection attempt" in msg] + assert len(retry_lines) == 2 + for message in retry_lines: + assert "testserver" not in message + assert "Trusted_Connection" not in message + + +def test_no_policy_adds_no_extra_log_lines(native, driver_log): + native.failures = 1 + with pytest.raises(OperationalError): + connect(CONN_STR) + assert [msg for level, msg in driver_log.messages if level == logging.WARNING] == [] + # Only the one error line _raise_connection_error has always written. + errors = [msg for level, msg in driver_log.messages if level == logging.ERROR] + assert len(errors) == 1 + assert "Connection attempt" not in errors[0] + + +def test_retry_policy_is_exported_from_the_package(): + assert mssql_python.RetryPolicy is RetryPolicy + assert "RetryPolicy" in mssql_python.__all__