feat(postgres): allow password to be a sync/async callable

CHANGELOG.rst

@@ -8,6 +8,15 @@ Changelog
 1.1
 ===
 
+1.1.9
+-----
+
+Added
+^^^^^
+- PostgreSQL ``password`` credential now accepts a sync or async callable, resolved once per
+  new connection, to support short-lived credentials such as AWS RDS/Aurora IAM tokens and
+  Azure Entra ID tokens. (#2034)
+
 1.1.8
 -----
 

docs/databases.rst

@@ -169,6 +169,48 @@ PostgreSQL optional parameters are pass-though parameters to the driver, see `he
 
 In case any of ``user``, ``password``, ``host``, ``port`` parameters is missing, we are letting ``asyncpg``/``psycopg`` retrieve it from default sources (standard PostgreSQL environment variables or default values).
 
+.. _db_password_callable:
+
+Rotating credentials
+--------------------
+
+``password`` also accepts a callable, which is invoked every time a new connection is
+opened. The callable may be synchronous or asynchronous, and must return a string.
+
+This makes Tortoise usable with short-lived credentials such as AWS RDS/Aurora IAM
+authentication tokens, Azure Entra ID tokens or Vault leases, where the password expires
+long before the pool does:
+
+.. code-block:: python3
+
+    async def get_token() -> str:
+        return await mint_short_lived_token()
+
+    await Tortoise.init(
+        config={
+            "connections": {
+                "default": {
+                    "engine": "tortoise.backends.asyncpg",
+                    "credentials": {
+                        "host": "db.host",
+                        "port": 5432,
+                        "user": "someuser",
+                        "password": get_token,
+                        "database": "somedb",
+                    },
+                }
+            },
+            "apps": {...},
+        }
+    )
+
+The callable is responsible for caching: it is awaited for every new connection the pool
+creates, so an expensive token request should be memoized until shortly before expiry.
+
+.. note::
+   Callable passwords cannot be expressed in a DB URL, so this requires dictionary
+   configuration. They are supported on the ``asyncpg`` and ``psycopg`` backends only.
+
 
 MySQL/MariaDB
 =============

tests/backends/test_password_factory.py

@@ -0,0 +1,144 @@
+from typing import Any
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from tortoise import Tortoise
+from tortoise.backends.base.client import resolve_password
+from tortoise.context import TortoiseContext
+
+
+@pytest.mark.asyncio
+async def test_resolve_password_passes_through_plain_values():
+    assert await resolve_password("foomip") == "foomip"
+    assert await resolve_password(None) is None
+
+
+@pytest.mark.asyncio
+async def test_resolve_password_calls_sync_and_async_callables():
+    calls = []
+
+    def sync_factory() -> str:
+        calls.append("sync")
+        return "sync-token"
+
+    async def async_factory() -> str:
+        calls.append("async")
+        return "async-token"
+
+    assert await resolve_password(sync_factory) == "sync-token"
+    assert await resolve_password(async_factory) == "async-token"
+    assert calls == ["sync", "async"]
+
+
+@pytest.mark.asyncio
+async def test_asyncpg_forwards_password_callable_to_the_driver():
+    """asyncpg resolves the callable itself, once per new connection."""
+    try:
+        import asyncpg  # noqa: F401
+    except ImportError:
+        pytest.skip("asyncpg not installed")
+
+    async def token() -> str:
+        return "token"
+
+    with patch(
+        "tortoise.backends.asyncpg.client.asyncpg.create_pool", new=AsyncMock()
+    ) as asyncpg_connect:
+        ctx = TortoiseContext()
+        async with ctx:
+            await ctx.connections._init(
+                {
+                    "models": {
+                        "engine": "tortoise.backends.asyncpg",
+                        "credentials": {
+                            "database": "test",
+                            "host": "127.0.0.1",
+                            "password": token,
+                            "port": 5432,
+                            "user": "root",
+                        },
+                    }
+                },
+                False,
+            )
+            await ctx.connections.get("models").create_connection(with_db=True)
+
+            assert asyncpg_connect.await_args.kwargs["password"] is token
+
+
+@pytest.mark.asyncio
+async def test_psycopg_keeps_the_password_out_of_the_conninfo():
+    try:
+        import psycopg  # noqa: F401
+    except ImportError:
+        pytest.skip("psycopg not installed")
+
+    def token() -> str:
+        return "token"
+
+    with patch(
+        "tortoise.backends.psycopg.client.PsycopgClient.create_pool", new=AsyncMock()
+    ) as patched_create_pool:
+        patched_create_pool.return_value = AsyncMock()
+        ctx = TortoiseContext()
+        async with ctx:
+            await ctx.connections._init(
+                {
+                    "models": {
+                        "engine": "tortoise.backends.psycopg",
+                        "credentials": {
+                            "database": "test",
+                            "host": "127.0.0.1",
+                            "password": token,
+                            "port": 5432,
+                            "user": "root",
+                            "timeout": 1,
+                        },
+                    }
+                },
+                False,
+            )
+            client = ctx.connections.get("models")
+            await client.create_connection(with_db=True)
+
+            assert "password" not in client._template["conninfo"]
+            assert client._template["connection_class"].__name__ == "PasswordFactoryConnection"
+
+
+@pytest.mark.asyncio
+async def test_psycopg_connection_class_mints_a_password_per_connection():
+    try:
+        from tortoise.backends.psycopg.client import password_factory_connection_class
+    except ImportError:
+        pytest.skip("psycopg not installed")
+
+    tokens = iter(["token-1", "token-2"])
+    seen: list[dict[str, Any]] = []
+
+    class RecordingConnection:
+        @classmethod
+        async def connect(cls, conninfo: str = "", **kwargs: Any) -> str:
+            seen.append(kwargs)
+            return "connection"
+
+    async def token() -> str:
+        return next(tokens)
+
+    connection_class = password_factory_connection_class(RecordingConnection, token)  # type: ignore[arg-type]
+
+    assert await connection_class.connect("host=127.0.0.1") == "connection"
+    assert await connection_class.connect("host=127.0.0.1") == "connection"
+    assert [kwargs["password"] for kwargs in seen] == ["token-1", "token-2"]
+
+
+def test_star_password_ignores_password_callables():
+    def factory() -> str:
+        return "s3cret"  # pragma: nocoverage
+
+    config = {
+        "models": {"engine": "tortoise.backends.asyncpg", "credentials": {"password": factory}}
+    }
+
+    # Must not raise: masking only applies to string passwords
+    assert "s3cret" not in Tortoise.star_password(config)

tortoise/__init__.py

@@ -433,7 +433,9 @@ def star_password(connections_config) -> str:
         for name, info in connections_config.items():
             if isinstance(info, str):
                 info = expand_db_url(info)
-            if password := info.get("credentials", {}).get("password"):
+            if (password := info.get("credentials", {}).get("password")) and isinstance(
+                password, str
+            ):
                 passwords.append(password)
 
         str_connection_config = str(connections_config)

tortoise/backends/asyncpg/client.py

@@ -58,6 +58,7 @@ async def create_connection(self, with_db: bool) -> None:
             **self.extra,
         }
         try:
+            # asyncpg resolves a callable password itself, once per new connection
             self._pool = await self.create_pool(password=self.password, **self._template)
             await self._post_connect()
             self.log.debug("Created connection pool %s with params: %s", self._pool, self._template)

tortoise/backends/base/client.py

@@ -2,8 +2,9 @@
 
 import abc
 import asyncio
-from collections.abc import Sequence
-from typing import Any, Generic, TypeVar, cast
+import inspect
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Any, Generic, TypeAlias, TypeVar, cast
 
 from pypika_tortoise import Query
 
@@ -15,6 +16,28 @@
 
 T_conn = TypeVar("T_conn")  # Instance of client connection, such as: asyncpg.Connection()
 
+#: A callable returning a password, awaited first if it returns an awaitable.
+PasswordFactory: TypeAlias = Callable[[], "str | Awaitable[str]"]
+PasswordType: TypeAlias = "str | PasswordFactory | None"
+
+
+async def resolve_password(password: PasswordType) -> str | None:
+    """
+    Resolve a password that may be supplied as a plain string or as a callable.
+
+    Callables are invoked once per connection attempt, which allows short-lived
+    credentials (AWS RDS IAM tokens, Azure Entra ID tokens, Vault leases, ...) to be
+    refreshed transparently without recreating the client.
+
+    :param password: A string, ``None``, or a sync/async callable returning a string.
+    """
+    if callable(password):
+        resolved = password()
+        if inspect.isawaitable(resolved):
+            resolved = await resolved
+        return resolved
+    return password
+
 
 class Capabilities:
     """

tortoise/backends/base_postgres/client.py

@@ -13,6 +13,7 @@
     BaseDBAsyncClient,
     Capabilities,
     ConnectionWrapper,
+    PasswordType,
     PoolConnectionWrapper,
     TransactionContext,
 )
@@ -61,7 +62,7 @@ class BasePostgresClient(BaseDBAsyncClient, abc.ABC):
     def __init__(
         self,
         user: str | None = None,
-        password: str | None = None,
+        password: PasswordType = None,
         database: str | None = None,
         host: str | None = None,
         port: SupportsInt = 5432,

tortoise/backends/psycopg/client.py

@@ -34,6 +34,25 @@ async def release(self, connection: psycopg.AsyncConnection):
         await self.putconn(connection)
 
 
+def password_factory_connection_class(
+    base: type[psycopg.AsyncConnection], factory: base_client.PasswordFactory
+) -> type[psycopg.AsyncConnection]:
+    """
+    Build a connection class that mints a fresh password for every new connection.
+
+    ``psycopg`` bakes credentials into an immutable conninfo string, so a rotating
+    password has to be injected at connect time instead.
+    """
+
+    class PasswordFactoryConnection(base):  # type: ignore[valid-type,misc]
+        @classmethod
+        async def connect(cls, conninfo: str = "", **kwargs: Any) -> Any:
+            kwargs["password"] = await base_client.resolve_password(factory)
+            return await super().connect(conninfo, **kwargs)
+
+    return PasswordFactoryConnection
+
+
 class PsycopgSQLQuery(PostgreSQLQuery):
     @classmethod
     def _builder(cls, **kwargs) -> PostgreSQLQueryBuilder:
@@ -82,11 +101,15 @@ async def create_connection(self, with_db: bool) -> None:
             host=self.host,
             port=self.port,
             user=self.user,
-            password=self.password,
+            password=None if callable(self.password) else self.password,
             dbname=self.database if with_db else None,
             **self.server_settings,
         )
 
+        connection_class: type[psycopg.AsyncConnection] = psycopg.AsyncConnection
+        if callable(self.password):
+            connection_class = password_factory_connection_class(connection_class, self.password)
+
         self._template = {
             "conninfo": conninfo,
             "min_size": self.pool_minsize,
@@ -95,7 +118,7 @@ async def create_connection(self, with_db: bool) -> None:
                 "autocommit": True,
                 "row_factory": psycopg.rows.dict_row,
             },
-            "connection_class": psycopg.AsyncConnection,
+            "connection_class": connection_class,
             **extra,
         }