diff --git a/CHANGELOG.md b/CHANGELOG.md index c4a2c48df..59ad247e4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,13 @@ 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 with a delay cap and full jitter + (on by default; `jitter=False` gives exact delays). `max_attempts` counts + total tries including the first; without a policy + `connect()` behaves exactly as before. ### Changed - `mssql-python` now depends on `mssql-python-rs==0.1.0` for `mssql_py_core` diff --git a/mssql_python/__init__.py b/mssql_python/__init__.py index e859d77c0..bba064b81 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 @@ -355,6 +358,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 21f564add..ae433391c 100644 --- a/mssql_python/connection.py +++ b/mssql_python/connection.py @@ -29,6 +29,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 ( @@ -246,6 +247,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: @@ -393,6 +414,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: """ @@ -459,6 +481,20 @@ 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, and the final failure after a retry, is logged at + warning level through the driver logger, which shows these lines once + ``setup_logging()`` has been called. **kwargs: Additional key/value pairs for the connection string. Returns: @@ -469,6 +505,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, @@ -490,6 +527,15 @@ 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. + 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 ) @@ -856,16 +902,63 @@ 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 + # attempt > 1 means at least one retry already happened. Without a policy + # max_attempts is 1, so a failure on the first try logs only the usual error line. + if attempt > 1: + logger.warning( + "Connection failed on attempt %d of %d with SQLSTATE %s; not retrying", + attempt, + max_attempts, + sqlstate or "none", + ) + _raise_connection_error(e) + except Exception: # pylint: disable=broad-exception-caught + # Anything other than a native connect error, such as one raised by a deferred + # token factory, is never retried and keeps its own type. After a retry the + # attempt that gave up is still logged. + if attempt > 1: + logger.warning( + "Connection failed on attempt %d of %d with SQLSTATE %s; not retrying", + attempt, + max_attempts, + "none", + ) + raise 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 ec7067093..73772cffa 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 81222f163..7f6681932 100644 --- a/mssql_python/mssql_python.pyi +++ b/mssql_python/mssql_python.pyi @@ -273,6 +273,11 @@ class _ArrowReader: use_internal_transaction: bool = False, ) -> Dict[str, Any]: ... +# Types used by the connect() / Connection signatures below, re-exported from the +# annotated implementations so they stay the single source of truth. +from .retry import RetryPolicy as RetryPolicy +from .connection import TokenProvider as TokenProvider + # DB-API 2.0 Connection Object # https://www.python.org/dev/peps/pep-0249/#connection-objects class Connection: @@ -312,6 +317,8 @@ class Connection: attrs_before: Optional[Dict[int, Union[int, str, bytes]]] = None, timeout: int = 0, native_uuid: Optional[bool] = None, + token_provider: Optional[TokenProvider] = None, + retry_policy: Optional[RetryPolicy] = None, **kwargs: Any, ) -> None: ... @@ -357,6 +364,8 @@ def connect( attrs_before: Optional[Dict[int, Union[int, str, bytes]]] = None, timeout: int = 0, native_uuid: Optional[bool] = None, + token_provider: Optional[TokenProvider] = 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 000000000..26f191733 --- /dev/null +++ b/mssql_python/retry.py @@ -0,0 +1,253 @@ +""" +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 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 + +# Upper bound for base_delay and max_delay: one day. It keeps every wait well inside the range +# time.sleep accepts on all supported platforms (about 49.7 days on Windows and about 3.2 years +# on macOS under Python 3.10), so a policy that validates can never fail inside the retry loop. +_MAX_DELAY_SECONDS = 86400.0 + + +def _is_valid_delay(value: object) -> bool: + """Return True for an int or float, not a bool, from zero to ``_MAX_DELAY_SECONDS``. + + NaN and infinity fail the range comparison, and an int too large for a float is compared + exactly, so neither needs a separate check. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return 0 <= value <= _MAX_DELAY_SECONDS + + +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 not iterable, or any code is not a + string of exactly five ASCII letters or digits. + """ + 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" + ) + # Only iter() sits in the try, so a TypeError raised inside a caller's generator is not + # relabelled as a bad setting. + try: + iterator = iter(codes) + except TypeError: + raise ValueError( + "retriable_sqlstates must be an iterable of SQLSTATE strings, " + f"got {type(codes).__name__}" + ) from None + normalized = set() + for code in iterator: + # Checked on the original string: ASCII only means upper() cannot change the length, + # and letters or digits are the only characters a driver SQLSTATE is parsed from. + if ( + not isinstance(code, str) + or len(code) != _SQLSTATE_LENGTH + or not (code.isascii() and code.isalnum()) + ): + raise ValueError( + f"each SQLSTATE must be {_SQLSTATE_LENGTH} ASCII letters or digits, 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, before jitter. + 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, 1), so many clients do not reconnect in lockstep. The delay can be shorter than + ``base_delay`` and can be zero. + retriable_sqlstates (frozenset): The SQLSTATE codes that are retried, uppercased and + each exactly five ASCII letters or digits. 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, before jitter; + zero to 86400 (one day). + max_delay (float): Cap in seconds for every delay; at least ``base_delay`` and at + most 86400. + jitter (bool): Scale each delay down by a random factor in [0, 1), so the wait can + be anywhere between zero and the backoff delay. + 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_valid_delay(base_delay): + raise ValueError("base_delay must be a number of seconds from 0 to 86400") + if not _is_valid_delay(max_delay) or max_delay < base_delay: + raise ValueError("max_delay must be a number of seconds from base_delay to 86400") + 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, before jitter.""" + 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 down by a random factor in [0, 1).""" + 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``. With jitter + on, the capped delay is scaled by a random factor in [0, 1), so zero is possible. + + 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: + # Full jitter: scale down by a factor in [0, 1) rather than around the delay. Scaling + # around it meant every draw at or above the midpoint clamped to max_delay, so once + # backoff reached the cap about half of all clients waited the identical amount. + delay *= _random() + 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 000000000..2e38fbec5 --- /dev/null +++ b/tests/test_027_retry_policy.py @@ -0,0 +1,520 @@ +""" +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 gc +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 InterfaceError, 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) + + +@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(caplog): + """Capture what the driver logger emits for the duration of a test. + + The driver logger does not propagate, so caplog's handler goes on it directly. + """ + # A failed connect() leaves a half built Connection in a reference cycle; if the collector + # frees one from an earlier test inside this window, its cleanup warning would be counted. + gc.collect() + with caplog.at_level(logging.WARNING, logger="mssql_python"): + mssql_python.logging.logger.addHandler(caplog.handler) + try: + yield caplog + finally: + mssql_python.logging.logger.removeHandler(caplog.handler) + + +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 exactly the same arguments. + 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_non_runtime_error_from_native_is_not_retried_or_rewrapped(monkeypatch, sleeps): + # An exception raised inside the deferred token factory reaches Python as its own type, + # not as a RuntimeError, so the retry loop must let it through untouched. + def fail(*args): + fail.calls += 1 + raise InterfaceError(driver_error="token factory failed", ddbc_error="") + + fail.calls = 0 + monkeypatch.setattr(mssql_python.connection.ddbc_bindings, "Connection", fail) + with pytest.raises(InterfaceError) as exc_info: + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert type(exc_info.value) is InterfaceError + assert fail.calls == 1 + assert sleeps == [] + + +def test_non_runtime_error_after_a_retry_logs_the_give_up_and_keeps_its_type( + monkeypatch, sleeps, driver_log +): + # A deferred token factory can fail with its own exception type on a later attempt. + def fail(*args): + fail.calls += 1 + if fail.calls == 1: + raise RuntimeError(LINK_FAILURE) + raise InterfaceError(driver_error="token factory failed", ddbc_error="") + + fail.calls = 0 + monkeypatch.setattr(mssql_python.connection.ddbc_bindings, "Connection", fail) + with pytest.raises(InterfaceError) as exc_info: + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert type(exc_info.value) is InterfaceError + assert fail.calls == 2 + assert sleeps == [1.0] + lines = [r.getMessage() for r in driver_log.records if "attempt" in r.getMessage()] + assert len(lines) == 2 + assert "attempt 1 of 3" in lines[0] and "08S01" in lines[0] + assert "attempt 2 of 3" in lines[1] and "SQLSTATE none;" in lines[1] + + +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_down_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.0, 0.0, 0.0] + monkeypatch.setattr(mssql_python.retry, "_random", lambda: 0.5) + assert [policy.compute_delay(n) for n in (1, 2, 3, 4)] == [0.5, 1.0, 2.0, 2.5] + + +def test_jitter_keeps_capped_delays_spread_out(monkeypatch): + # Once backoff reaches max_delay every client is asking for the same number, so the jitter is + # the only thing keeping them apart. Scaling around the delay used to clamp roughly half of + # the draws to exactly max_delay. + # The seam is fed an evenly spaced sweep of [0, 1) rather than a seeded generator. It covers + # the interval the same way, so the assertions below hold on every run instead of for one + # seed, and the test carries no random source of its own. + sweep = iter([(i + 0.5) / 2000 for i in range(2000)]) + monkeypatch.setattr(mssql_python.retry, "_random", lambda: next(sweep)) + policy = RetryPolicy(base_delay=1.0, max_delay=30.0, jitter=True) + delays = [policy.compute_delay(5000) for _ in range(2000)] + assert all(0.0 <= d < 30.0 for d in delays) + # a uniform draw over [0, 30) should not pile up in any one tenth of the range + buckets = [0] * 10 + for d in delays: + buckets[int(d / 3.0)] += 1 + assert max(buckets) < len(delays) / 4 + + +def test_jittered_delays_are_used_when_retrying(native, sleeps, monkeypatch): + monkeypatch.setattr(mssql_python.retry, "_random", lambda: 0.5) + 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({"base_delay": 10**400}, id="base_delay_huge_int"), + pytest.param({"max_delay": 10**400}, id="max_delay_huge_int"), + pytest.param({"max_delay": 86400.5}, id="max_delay_above_limit"), + 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"), + pytest.param({"retriable_sqlstates": 123}, id="sqlstate_not_iterable"), + pytest.param({"retriable_sqlstates": ["08-01"]}, id="sqlstate_bad_chars"), + pytest.param({"retriable_sqlstates": ["08ßAB"]}, id="sqlstate_non_ascii"), + ], +) +def test_invalid_settings_raise_value_error(kwargs): + with pytest.raises(ValueError): + RetryPolicy(**kwargs) + + +def test_delays_up_to_one_day_are_accepted(): + policy = RetryPolicy(base_delay=86400, max_delay=86400, jitter=False) + assert policy.compute_delay(1) == 86400.0 + + +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_wrong_policy_type_fails_before_a_token_is_acquired(native, sleeps): + provider = CountingTokenProvider() + with pytest.raises(TypeError): + connect("Server=testserver;Database=mydb;", token_provider=provider, retry_policy="nope") + assert provider.calls == 0 + assert native.calls == [] + + +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] + # Here attrs_before carries the token and the pool key is not empty, so this also shows + # that neither changes between attempts. + assert all(call == native.calls[0] for call in native.calls) + + +def test_deferred_token_factory_and_pool_key_are_reused_on_every_attempt(native, sleeps): + # ActiveDirectoryMsi builds a token factory that the native layer calls on each physical + # connect. Building it does no network work, and the fake native never calls it. + native.failures = 2 + connect( + "Server=testserver;Database=mydb;Authentication=ActiveDirectoryMsi;", + retry_policy=RetryPolicy(max_attempts=3, jitter=False), + ) + assert len(native.calls) == 3 + assert callable(native.calls[0][4]) + assert len({id(call[4]) for call in native.calls}) == 1 + assert native.calls[0][3] + assert all(call[3] == native.calls[0][3] for call in native.calls) + + +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 = [r.getMessage() for r in driver_log.records if r.levelno == logging.WARNING] + errors = [r.getMessage() for r in driver_log.records if r.levelno == logging.ERROR] + assert len(warnings) == 3 + 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] + # Giving up adds one warning that says how many attempts ran. + assert "3 of 3" in warnings[2] and "08S01" in warnings[2] and "not retrying" in warnings[2] + # The final failure still logs only the one error line _raise_connection_error has always + # written. + assert len(errors) == 1 + assert "Connection attempt" not in errors[0] + for message in warnings: + assert "testserver" not in message + assert "Trusted_Connection" not in message + + +@pytest.mark.parametrize( + "second, sqlstate", + [ + pytest.param(LOGIN_FAILURE, "28000", id="permanent_sqlstate"), + pytest.param(DRIVER_PREFIX + "Connection handle not allocated", "none", id="no_sqlstate"), + ], +) +def test_failure_after_a_retry_logs_the_attempt_that_gave_up( + monkeypatch, sleeps, driver_log, second, sqlstate +): + messages = [LINK_FAILURE, second] + + def fail(*args): + fail.calls += 1 + raise RuntimeError(messages[fail.calls - 1]) + + fail.calls = 0 + monkeypatch.setattr(mssql_python.connection.ddbc_bindings, "Connection", fail) + with pytest.raises(OperationalError): + connect(CONN_STR, retry_policy=RetryPolicy(max_attempts=3, jitter=False)) + assert fail.calls == 2 + assert sleeps == [1.0] + # Only the retry loop's own lines are counted. A Connection left half built by a failed + # connect logs a cleanup warning whenever the garbage collector reaches it, which can be now. + lines = [r.getMessage() for r in driver_log.records if "attempt" in r.getMessage()] + assert len(lines) == 2 + assert "attempt 1 of 3" in lines[0] and "08S01" in lines[0] + assert "attempt 2 of 3" in lines[1] and f"SQLSTATE {sqlstate};" in lines[1] + + +def test_no_policy_adds_no_extra_log_lines(native, driver_log): + native.failures = 1 + with pytest.raises(OperationalError): + connect(CONN_STR) + assert [r.getMessage() for r in driver_log.records if r.levelno == logging.WARNING] == [] + # Only the one error line _raise_connection_error has always written. + errors = [r.getMessage() for r in driver_log.records if r.levelno == 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__ diff --git a/tests/test_028_stub_signature_parity.py b/tests/test_028_stub_signature_parity.py new file mode 100644 index 000000000..00a1d1fe1 --- /dev/null +++ b/tests/test_028_stub_signature_parity.py @@ -0,0 +1,77 @@ +""" +Pins the type stub to the runtime signatures of connect() and Connection.__init__. + +Parses source with ast only, so it needs neither the native ddbc_bindings module nor a server. +Annotations are not compared because the stub spells them differently (a string forward +reference, an unqualified RetryPolicy); kind, name, order and default are. +""" + +import ast +from pathlib import Path + +import pytest + +PKG = Path(__file__).resolve().parent.parent / "mssql_python" + + +def _signature(path, name, cls=None): + """Return [(kind, name, default-source)] for a def in path, optionally inside class cls.""" + body = ast.parse(path.read_text(encoding="utf-8")).body + if cls is not None: + body = next(n for n in body if isinstance(n, ast.ClassDef) and n.name == cls).body + func = next(n for n in body if isinstance(n, ast.FunctionDef) and n.name == name) + a = func.args + positional = a.posonlyargs + a.args + # Tag the two positional groups apart. A single tag for both would let a stub that drops the + # slash compare equal to a runtime that keeps it. a.defaults spans posonlyargs and args + # jointly, so the defaults list is still built over the combined sequence. + kinds = ["posonly"] * len(a.posonlyargs) + ["pos"] * len(a.args) + defaults = [None] * (len(positional) - len(a.defaults)) + a.defaults + params = [ + (kind, arg.arg, ast.unparse(d) if d is not None else None) + for kind, arg, d in zip(kinds, positional, defaults) + ] + if a.vararg: + params.append(("*", a.vararg.arg, None)) + params += [ + ("kwonly", arg.arg, ast.unparse(d) if d is not None else None) + for arg, d in zip(a.kwonlyargs, a.kw_defaults) + ] + if a.kwarg: + params.append(("**", a.kwarg.arg, None)) + return params + + +@pytest.mark.parametrize( + "runtime_file, name, cls", + [ + ("connection.py", "__init__", "Connection"), + ("db_connection.py", "connect", None), + ], +) +def test_stub_matches_runtime_signature(runtime_file, name, cls): + runtime = _signature(PKG / runtime_file, name, cls) + stub = _signature(PKG / "mssql_python.pyi", name, cls) + assert stub == runtime + + +def test_connect_forwards_every_connection_parameter(): + init = _signature(PKG / "connection.py", "__init__", "Connection") + assert _signature(PKG / "db_connection.py", "connect") == init[1:] # drop self + + +def test_signature_separates_the_two_positional_kinds(tmp_path): + """A stub that drops the slash must not compare equal to a runtime that keeps it. + + Tagging posonlyargs and args alike would hide exactly the positional drift this file is + here to catch, so the kinds are pinned directly. + """ + runtime = tmp_path / "runtime.py" + runtime.write_text("def f(a, /, b=1): ...\n", encoding="utf-8") + stub = tmp_path / "stub.py" + stub.write_text("def f(a, b=1): ...\n", encoding="utf-8") + assert [kind for kind, _, _ in _signature(runtime, "f")] == ["posonly", "pos"] + assert [kind for kind, _, _ in _signature(stub, "f")] == ["pos", "pos"] + assert _signature(runtime, "f") != _signature(stub, "f") + # the default still has to travel with the parameter it belongs to + assert _signature(runtime, "f")[1] == ("pos", "b", "1")