Retry

langchain_cockroachdb.retry

Retry policy for CockroachDB operations.

CockroachDB runs transactions at SERIALIZABLE isolation by default, so clients are expected to retry some errors. This module decides which errors are worth retrying and runs operations with exponential backoff and jitter.

A 40001 error means the transaction did not commit, so retrying it is always safe. After 40003 (result is ambiguous) or a lost connection the transaction may have committed, so only idempotent work should be retried. Every write this package issues is idempotent (UPSERT, ON CONFLICT DO NOTHING, DELETE, IF [NOT] EXISTS).

RetryPolicy dataclass

Which CockroachDB errors to retry, and how.

Every field can be changed. RetryPolicy(**mapping) accepts plain lists for the code and phrase fields, so policies can come from YAML or environment variables. Use dataclasses.replace to derive a policy.

Attributes:

Name	Type	Description
max_attempts	int	Total attempts, including the first (default: 5)
initial_backoff	float	Delay before the first retry, in seconds (default: 0.1)
max_backoff	float	Upper bound for any delay, in seconds (default: 10.0)
backoff_multiplier	float	Growth factor between delays (default: 2.0)
jitter	float	Fraction of each delay that is randomized, 0 to 1 (default: 0.5)
retryable_sqlstates	frozenset[str]	SQLSTATEs to retry (default: RETRYABLE_SQLSTATES)
retryable_sqlstate_classes	frozenset[str]	Two-character SQLSTATE classes to retry (default: RETRYABLE_SQLSTATE_CLASSES)
uncategorized_sqlstates	frozenset[str]	Codes too generic to decide on; message phrases decide instead (default: UNCATEGORIZED_SQLSTATES)
transient_messages	tuple[str, ...]	Phrases that mark an error without a usable code as transient (default: TRANSIENT_MESSAGES)
non_transient_messages	tuple[str, ...]	Phrases that rule out a retry for errors without a usable code (default: NON_TRANSIENT_MESSAGES)
retry_on	Callable[[BaseException], bool] | None	Custom predicate that replaces the rules above
attempt_timeout	float | None	Seconds before a hung query or write is abandoned and retried; None disables the guard (default: None)
ddl_attempt_timeout	float | None	Same guard, bounds each schema change (default: None)

Example

# Ride through multi-region failovers: also retry statement timeouts
# and cut off attempts that hang for more than 10 seconds.
policy = RetryPolicy(
    retryable_sqlstates=RETRYABLE_SQLSTATES | {"57014"},
    attempt_timeout=10,
)

Source code in langchain_cockroachdb/retry.py

@dataclass(frozen=True)
class RetryPolicy:
    """Which CockroachDB errors to retry, and how.

    Every field can be changed. ``RetryPolicy(**mapping)`` accepts plain lists
    for the code and phrase fields, so policies can come from YAML or
    environment variables. Use ``dataclasses.replace`` to derive a policy.

    Attributes:
        max_attempts: Total attempts, including the first (default: 5)
        initial_backoff: Delay before the first retry, in seconds (default: 0.1)
        max_backoff: Upper bound for any delay, in seconds (default: 10.0)
        backoff_multiplier: Growth factor between delays (default: 2.0)
        jitter: Fraction of each delay that is randomized, 0 to 1 (default: 0.5)
        retryable_sqlstates: SQLSTATEs to retry (default: RETRYABLE_SQLSTATES)
        retryable_sqlstate_classes: Two-character SQLSTATE classes to retry
            (default: RETRYABLE_SQLSTATE_CLASSES)
        uncategorized_sqlstates: Codes too generic to decide on; message
            phrases decide instead (default: UNCATEGORIZED_SQLSTATES)
        transient_messages: Phrases that mark an error without a usable code
            as transient (default: TRANSIENT_MESSAGES)
        non_transient_messages: Phrases that rule out a retry for errors
            without a usable code (default: NON_TRANSIENT_MESSAGES)
        retry_on: Custom predicate that replaces the rules above
        attempt_timeout: Seconds before a hung query or write is abandoned and
            retried; None disables the guard (default: None)
        ddl_attempt_timeout: Same guard, bounds each schema change (default: None)

    Example:
        ```python
        # Ride through multi-region failovers: also retry statement timeouts
        # and cut off attempts that hang for more than 10 seconds.
        policy = RetryPolicy(
            retryable_sqlstates=RETRYABLE_SQLSTATES | {"57014"},
            attempt_timeout=10,
        )
        ```
    """

    max_attempts: int = 5
    initial_backoff: float = 0.1
    max_backoff: float = 10.0
    backoff_multiplier: float = 2.0
    jitter: float = 0.5
    retryable_sqlstates: frozenset[str] = RETRYABLE_SQLSTATES
    retryable_sqlstate_classes: frozenset[str] = RETRYABLE_SQLSTATE_CLASSES
    uncategorized_sqlstates: frozenset[str] = UNCATEGORIZED_SQLSTATES
    transient_messages: tuple[str, ...] = TRANSIENT_MESSAGES
    non_transient_messages: tuple[str, ...] = NON_TRANSIENT_MESSAGES
    retry_on: Callable[[BaseException], bool] | None = None
    attempt_timeout: float | None = None
    ddl_attempt_timeout: float | None = None

    def __post_init__(self) -> None:
        if (
            isinstance(self.max_attempts, bool)
            or not isinstance(self.max_attempts, int)
            or self.max_attempts < 1
        ):
            raise ValueError("max_attempts must be an integer of at least 1")
        if self.initial_backoff < 0:
            raise ValueError("initial_backoff must not be negative")
        if self.max_backoff < self.initial_backoff:
            raise ValueError("max_backoff must be at least initial_backoff")
        if self.backoff_multiplier < 1:
            raise ValueError("backoff_multiplier must be at least 1")
        if isinstance(self.jitter, bool):
            raise TypeError("jitter is a ratio from 0 to 1, for example 0.5 or 0")
        if not 0 <= self.jitter <= 1:
            raise ValueError("jitter must be between 0 and 1")
        for name in ("attempt_timeout", "ddl_attempt_timeout"):
            value = getattr(self, name)
            if value is not None and value <= 0:
                raise ValueError(f"{name} must be positive or None")
        for name in (
            "retryable_sqlstates",
            "retryable_sqlstate_classes",
            "uncategorized_sqlstates",
        ):
            object.__setattr__(self, name, _as_codes(getattr(self, name), name))
        for name in ("transient_messages", "non_transient_messages"):
            object.__setattr__(self, name, _as_phrases(getattr(self, name), name))

    def is_retryable(self, error: BaseException) -> bool:
        """Return True if ``error`` is worth retrying under this policy."""
        if self.retry_on is not None:
            return bool(self.retry_on(error))
        code = get_sqlstate(error)
        if code is not None:
            if code in self.retryable_sqlstates or code[:2] in self.retryable_sqlstate_classes:
                return True
            if code not in self.uncategorized_sqlstates:
                return False
        chain = list(_error_chain(error))
        if not any(isinstance(e, _driver_error_types()) for e in chain):
            return False
        if any(isinstance(e, (PoolTimeout, PoolCheckoutTimeout)) for e in chain):
            return False
        # SQLAlchemy's message embeds the statement and its parameters, which
        # can hold user data, so only driver and application messages count.
        messages = [str(e).lower() for e in chain if not isinstance(e, DBAPIError)]
        if any(phrase in m for m in messages for phrase in self.non_transient_messages):
            return False
        if any(getattr(e, "connection_invalidated", False) for e in chain):
            return True
        return any(phrase in m for m in messages for phrase in self.transient_messages)

    def delay(self, retry_number: int, *, rand: Callable[[], float] = random.random) -> float:
        """Seconds to wait before retry ``retry_number`` (0 for the first retry)."""
        exponent = min(max(retry_number, 0), 64)
        base = min(self.initial_backoff * self.backoff_multiplier**exponent, self.max_backoff)
        return base * (1 - self.jitter * rand())

    def run(
        self,
        fn: Callable[[], T],
        *,
        label: str = "operation",
        can_retry: Callable[[BaseException], bool] | None = None,
        sleep: Callable[[float], Any] = time.sleep,
    ) -> T:
        """Call ``fn`` until it succeeds or fails with an error not worth retrying.

        Args:
            fn: The attempt; called again from scratch on every retry
            label: What the attempt does, used in log messages
            can_retry: Extra check; when it returns False the error is raised
            sleep: Sleep function (tests pass a fake)

        Returns:
            Whatever ``fn`` returns
        """
        attempt = 1
        while True:
            try:
                return fn()
            except Exception as error:
                if not self._should_retry(error, attempt, label, can_retry):
                    raise
                pause = self.delay(attempt - 1)
                self._log_retry(error, attempt, pause, label)
                sleep(pause)
                attempt += 1

    async def arun(
        self,
        fn: Callable[[], Awaitable[T]],
        *,
        label: str = "operation",
        can_retry: Callable[[BaseException], bool] | None = None,
        sleep: Callable[[float], Awaitable[Any]] = asyncio.sleep,
    ) -> T:
        """Async version of :meth:`run`; ``fn`` returns a new awaitable per attempt."""
        attempt = 1
        while True:
            try:
                return await fn()
            except Exception as error:
                if not self._should_retry(error, attempt, label, can_retry):
                    raise
                pause = self.delay(attempt - 1)
                self._log_retry(error, attempt, pause, label)
                await sleep(pause)
                attempt += 1

    def _should_retry(
        self,
        error: BaseException,
        attempt: int,
        label: str,
        can_retry: Callable[[BaseException], bool] | None,
    ) -> bool:
        if not self.is_retryable(error):
            return False
        if can_retry is not None and not can_retry(error):
            return False
        if attempt >= self.max_attempts:
            logger.warning(
                "Giving up on %s after %d attempts: %s", label, attempt, _describe(error)
            )
            return False
        return True

    def _log_retry(self, error: BaseException, attempt: int, pause: float, label: str) -> None:
        logger.warning(
            "Retrying %s after %s (attempt %d of %d failed, next try in %.2fs)",
            label,
            _describe(error),
            attempt,
            self.max_attempts,
            pause,
        )

arun(fn, *, label='operation', can_retry=None, sleep=asyncio.sleep) async

Async version of :meth:run; fn returns a new awaitable per attempt.

Source code in langchain_cockroachdb/retry.py

async def arun(
    self,
    fn: Callable[[], Awaitable[T]],
    *,
    label: str = "operation",
    can_retry: Callable[[BaseException], bool] | None = None,
    sleep: Callable[[float], Awaitable[Any]] = asyncio.sleep,
) -> T:
    """Async version of :meth:`run`; ``fn`` returns a new awaitable per attempt."""
    attempt = 1
    while True:
        try:
            return await fn()
        except Exception as error:
            if not self._should_retry(error, attempt, label, can_retry):
                raise
            pause = self.delay(attempt - 1)
            self._log_retry(error, attempt, pause, label)
            await sleep(pause)
            attempt += 1

delay(retry_number, *, rand=random.random)

Seconds to wait before retry retry_number (0 for the first retry).

Source code in langchain_cockroachdb/retry.py

def delay(self, retry_number: int, *, rand: Callable[[], float] = random.random) -> float:
    """Seconds to wait before retry ``retry_number`` (0 for the first retry)."""
    exponent = min(max(retry_number, 0), 64)
    base = min(self.initial_backoff * self.backoff_multiplier**exponent, self.max_backoff)
    return base * (1 - self.jitter * rand())

is_retryable(error)

Return True if error is worth retrying under this policy.

Source code in langchain_cockroachdb/retry.py

def is_retryable(self, error: BaseException) -> bool:
    """Return True if ``error`` is worth retrying under this policy."""
    if self.retry_on is not None:
        return bool(self.retry_on(error))
    code = get_sqlstate(error)
    if code is not None:
        if code in self.retryable_sqlstates or code[:2] in self.retryable_sqlstate_classes:
            return True
        if code not in self.uncategorized_sqlstates:
            return False
    chain = list(_error_chain(error))
    if not any(isinstance(e, _driver_error_types()) for e in chain):
        return False
    if any(isinstance(e, (PoolTimeout, PoolCheckoutTimeout)) for e in chain):
        return False
    # SQLAlchemy's message embeds the statement and its parameters, which
    # can hold user data, so only driver and application messages count.
    messages = [str(e).lower() for e in chain if not isinstance(e, DBAPIError)]
    if any(phrase in m for m in messages for phrase in self.non_transient_messages):
        return False
    if any(getattr(e, "connection_invalidated", False) for e in chain):
        return True
    return any(phrase in m for m in messages for phrase in self.transient_messages)

run(fn, *, label='operation', can_retry=None, sleep=time.sleep)

Call fn until it succeeds or fails with an error not worth retrying.

Parameters:

Name	Type	Description	Default
fn	Callable[[], T]	The attempt; called again from scratch on every retry	required
label	str	What the attempt does, used in log messages	'operation'
can_retry	Callable[[BaseException], bool] | None	Extra check; when it returns False the error is raised	None
sleep	Callable[[float], Any]	Sleep function (tests pass a fake)	sleep

Returns:

Type	Description
T	Whatever fn returns

Source code in langchain_cockroachdb/retry.py

def run(
    self,
    fn: Callable[[], T],
    *,
    label: str = "operation",
    can_retry: Callable[[BaseException], bool] | None = None,
    sleep: Callable[[float], Any] = time.sleep,
) -> T:
    """Call ``fn`` until it succeeds or fails with an error not worth retrying.

    Args:
        fn: The attempt; called again from scratch on every retry
        label: What the attempt does, used in log messages
        can_retry: Extra check; when it returns False the error is raised
        sleep: Sleep function (tests pass a fake)

    Returns:
        Whatever ``fn`` returns
    """
    attempt = 1
    while True:
        try:
            return fn()
        except Exception as error:
            if not self._should_retry(error, attempt, label, can_retry):
                raise
            pause = self.delay(attempt - 1)
            self._log_retry(error, attempt, pause, label)
            sleep(pause)
            attempt += 1

async_retry_with_backoff(max_retries=None, initial_backoff=None, max_backoff=None, backoff_multiplier=None, jitter=None, *, retry_on=None, policy=None)

Retry an async function on transient CockroachDB errors.

Unset arguments fall back to RetryPolicy defaults (5 attempts, 0.1s initial backoff, 10s cap, doubling, jitter on). The whole function runs again on each attempt, so it should open its own transaction.

After 40003 or a lost connection the work may already have committed: only decorate idempotent work, or pass retry_on=lambda e: get_sqlstate(e) == "40001".

Parameters:

Name	Type	Description	Default
max_retries	int | None	Total attempts, including the first	None
initial_backoff	float | None	Delay before the first retry, in seconds	None
max_backoff	float | None	Upper bound for any delay, in seconds	None
backoff_multiplier	float | None	Growth factor between delays	None
jitter	bool | None	Randomize delays (True maps to a 0.5 ratio)	None
retry_on	Callable[[BaseException], bool] | None	Custom predicate deciding which errors to retry	None
policy	RetryPolicy | None	A complete RetryPolicy, instead of the arguments above	None

Example

@async_retry_with_backoff(max_retries=3)
async def upsert(conn):
    await conn.execute(text("UPSERT INTO t VALUES (1, 'a')"))

Source code in langchain_cockroachdb/retry.py

def async_retry_with_backoff(
    max_retries: int | None = None,
    initial_backoff: float | None = None,
    max_backoff: float | None = None,
    backoff_multiplier: float | None = None,
    jitter: bool | None = None,
    *,
    retry_on: Callable[[BaseException], bool] | None = None,
    policy: RetryPolicy | None = None,
) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
    """Retry an async function on transient CockroachDB errors.

    Unset arguments fall back to ``RetryPolicy`` defaults (5 attempts, 0.1s
    initial backoff, 10s cap, doubling, jitter on). The whole function runs
    again on each attempt, so it should open its own transaction.

    After 40003 or a lost connection the work may already have committed:
    only decorate idempotent work, or pass
    ``retry_on=lambda e: get_sqlstate(e) == "40001"``.

    Args:
        max_retries: Total attempts, including the first
        initial_backoff: Delay before the first retry, in seconds
        max_backoff: Upper bound for any delay, in seconds
        backoff_multiplier: Growth factor between delays
        jitter: Randomize delays (True maps to a 0.5 ratio)
        retry_on: Custom predicate deciding which errors to retry
        policy: A complete RetryPolicy, instead of the arguments above

    Example:
        ```python
        @async_retry_with_backoff(max_retries=3)
        async def upsert(conn):
            await conn.execute(text("UPSERT INTO t VALUES (1, 'a')"))
        ```
    """
    resolved = _decorator_policy(
        max_retries, initial_backoff, max_backoff, backoff_multiplier, jitter, retry_on, policy
    )

    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
        label = getattr(func, "__qualname__", None) or getattr(func, "__name__", None) or repr(func)

        @functools.wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> T:
            return await resolved.arun(lambda: func(*args, **kwargs), label=label)

        return wrapper

    return decorator

get_sqlstate(error)

Return the SQLSTATE carried by an error or by anything that caused it.

Reads sqlstate (psycopg 3, asyncpg) and pgcode (psycopg2), looking through SQLAlchemy's wrapped orig and raise ... from causes.

Parameters:

Name	Type	Description	Default
error	BaseException	The exception to inspect	required

Returns:

Type	Description
str | None	The five-character SQLSTATE, or None if nothing in the chain has one

Source code in langchain_cockroachdb/retry.py

def get_sqlstate(error: BaseException) -> str | None:
    """Return the SQLSTATE carried by an error or by anything that caused it.

    Reads ``sqlstate`` (psycopg 3, asyncpg) and ``pgcode`` (psycopg2), looking
    through SQLAlchemy's wrapped ``orig`` and ``raise ... from`` causes.

    Args:
        error: The exception to inspect

    Returns:
        The five-character SQLSTATE, or None if nothing in the chain has one
    """
    for current in _error_chain(error):
        for attr in ("sqlstate", "pgcode"):
            code = getattr(current, attr, None)
            if isinstance(code, str) and len(code) == 5:
                return code.upper()
    return None

is_retryable_error(error)

Return True if error is worth retrying under the default policy.

Source code in langchain_cockroachdb/retry.py

def is_retryable_error(error: BaseException) -> bool:
    """Return True if ``error`` is worth retrying under the default policy."""
    return DEFAULT_RETRY_POLICY.is_retryable(error)

resolve_retry_policy(policy, *, base, max_attempts=None, initial_backoff=None, max_backoff=None, backoff_multiplier=None, jitter=None)

Combine an explicit policy, or flat retry_* arguments, with a base policy.

Raises:

Type	Description
ValueError	If both a policy and flat arguments are given

Source code in langchain_cockroachdb/retry.py

def resolve_retry_policy(
    policy: RetryPolicy | None,
    *,
    base: RetryPolicy,
    max_attempts: int | None = None,
    initial_backoff: float | None = None,
    max_backoff: float | None = None,
    backoff_multiplier: float | None = None,
    jitter: bool | None = None,
) -> RetryPolicy:
    """Combine an explicit policy, or flat ``retry_*`` arguments, with a base policy.

    Raises:
        ValueError: If both a policy and flat arguments are given
    """
    overrides: dict[str, Any] = {
        "max_attempts": max_attempts,
        "initial_backoff": initial_backoff,
        "max_backoff": max_backoff,
        "backoff_multiplier": backoff_multiplier,
        "jitter": None if jitter is None else (0.5 if jitter else 0.0),
    }
    overrides = {name: value for name, value in overrides.items() if value is not None}
    if policy is not None:
        if overrides:
            raise ValueError("Pass either retry_policy or retry_* arguments, not both")
        return policy
    return dataclasses.replace(base, **overrides) if overrides else base

sync_retry_with_backoff(max_retries=None, initial_backoff=None, max_backoff=None, backoff_multiplier=None, jitter=None, *, retry_on=None, policy=None)

Retry a function on transient CockroachDB errors.

Synchronous version of :func:async_retry_with_backoff, with the same arguments and the same idempotency caveat.

Source code in langchain_cockroachdb/retry.py

def sync_retry_with_backoff(
    max_retries: int | None = None,
    initial_backoff: float | None = None,
    max_backoff: float | None = None,
    backoff_multiplier: float | None = None,
    jitter: bool | None = None,
    *,
    retry_on: Callable[[BaseException], bool] | None = None,
    policy: RetryPolicy | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Retry a function on transient CockroachDB errors.

    Synchronous version of :func:`async_retry_with_backoff`, with the same
    arguments and the same idempotency caveat.
    """
    resolved = _decorator_policy(
        max_retries, initial_backoff, max_backoff, backoff_multiplier, jitter, retry_on, policy
    )

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        label = getattr(func, "__qualname__", None) or getattr(func, "__name__", None) or repr(func)

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            return resolved.run(lambda: func(*args, **kwargs), label=label)

        return wrapper

    return decorator