Merge pull request #3 from community-of-python/update-and-rename

rename to db-try, sync with internal version

Author: lesnik512

.github/workflows/ci.yml

@@ -20,7 +20,7 @@ jobs:
         with:
           enable-cache: true
           cache-dependency-glob: "**/pyproject.toml"
-      - run: uv python install 3.10
+      - run: uv python install 3.13
       - run: just install lint-ci
 
   pytest:

Justfile

@@ -14,7 +14,7 @@ build:
 
 install:
     uv lock --upgrade
-    uv sync --all-extras --frozen
+    uv sync --all-extras --all-groups --frozen
 
 lint:
     uv run --frozen ruff format

README.md

@@ -1,6 +1,6 @@
-# modern-pg
+# db-try
 
 This library provides retry decorators for sqlalchemy and some helpers
 
-Default settings are in [./pg-tools/settings.py](modern_pg/settings.py).
+Default settings are in [./db_try/settings.py](db_try/settings.py).
 You can redefine them by environment variables.

db_try/__init__.py

@@ -0,0 +1,13 @@
+from db_try.connections import build_connection_factory
+from db_try.decorators import postgres_retry
+from db_try.helpers import build_db_dsn, is_dsn_multihost
+from db_try.transaction import Transaction
+
+
+__all__ = [
+    "Transaction",
+    "build_connection_factory",
+    "build_db_dsn",
+    "is_dsn_multihost",
+    "postgres_retry",
+]

modern_pg/connections.py renamed to db_try/connections.py

@@ -71,7 +71,7 @@ async def _connection_factory() -> "ConnectionType":
                     target_session_attrs=target_session_attrs,
                 )
                 return connection  # noqa: TRY300
-            except (TimeoutError, OSError, asyncpg.TargetServerAttributeNotMatched) as exc:  # noqa: PERF203
+            except (TimeoutError, OSError, asyncpg.TargetServerAttributeNotMatched) as exc:
                 logger.warning("Failed to fetch asyncpg connection from %s, %s", one_host, exc)
         msg: typing.Final = f"None of the hosts match the target attribute requirement {target_session_attrs}"
         raise asyncpg.TargetServerAttributeNotMatched(msg)

db_try/decorators.py

@@ -0,0 +1,42 @@
+import functools
+import logging
+import typing
+
+import asyncpg
+import tenacity
+from sqlalchemy.exc import DBAPIError
+
+from db_try import settings
+
+
+logger = logging.getLogger(__name__)
+
+
+def _retry_handler(exception: BaseException) -> bool:
+    if (
+        isinstance(exception, DBAPIError)
+        and hasattr(exception, "orig")
+        and isinstance(exception.orig.__cause__, (asyncpg.SerializationError, asyncpg.PostgresConnectionError))  # type: ignore[union-attr]
+    ):
+        logger.debug("postgres_retry, retrying")
+        return True
+
+    logger.debug("postgres_retry, giving up on retry")
+    return False
+
+
+def postgres_retry[**P, T](
+    func: typing.Callable[P, typing.Coroutine[None, None, T]],
+) -> typing.Callable[P, typing.Coroutine[None, None, T]]:
+    @tenacity.retry(
+        stop=tenacity.stop_after_attempt(settings.DB_UTILS_RETRIES_NUMBER),
+        wait=tenacity.wait_exponential_jitter(),
+        retry=tenacity.retry_if_exception(_retry_handler),
+        reraise=True,
+        before=tenacity.before_log(logger, logging.DEBUG),
+    )
+    @functools.wraps(func)
+    async def wrapped_method(*args: P.args, **kwargs: P.kwargs) -> T:
+        return await func(*args, **kwargs)
+
+    return wrapped_method

modern_pg/helpers.py renamed to db_try/helpers.py

@@ -0,0 +1,5 @@
+import os
+import typing
+
+
+DB_UTILS_RETRIES_NUMBER: typing.Final = int(os.getenv("DB_UTILS_RETRIES_NUMBER", "3"))

modern_pg/py.typed renamed to db_try/py.typed

@@ -0,0 +1,28 @@
+import dataclasses
+
+import typing_extensions
+from sqlalchemy.engine.interfaces import IsolationLevel
+from sqlalchemy.ext import asyncio as sa_async
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True, slots=True)
+class Transaction:
+    session: sa_async.AsyncSession
+    isolation_level: IsolationLevel | None = None
+
+    async def __aenter__(self) -> typing_extensions.Self:
+        if self.isolation_level:
+            await self.session.connection(execution_options={"isolation_level": self.isolation_level})
+
+        if not self.session.in_transaction():
+            await self.session.begin()
+        return self
+
+    async def __aexit__(self, *args: object, **kwargs: object) -> None:
+        await self.session.close()
+
+    async def commit(self) -> None:
+        await self.session.commit()
+
+    async def rollback(self) -> None:
+        await self.session.rollback()