diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9d9957975..5818eb6ae 100644 --- a/CHANGELOG.rst +++ b/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 ----- diff --git a/docs/databases.rst b/docs/databases.rst index 4341e7d2d..7685474e2 100644 --- a/docs/databases.rst +++ b/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 ============= diff --git a/tests/backends/test_password_factory.py b/tests/backends/test_password_factory.py new file mode 100644 index 000000000..29ff9c29d --- /dev/null +++ b/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) diff --git a/tortoise/__init__.py b/tortoise/__init__.py index 34dc7e133..4ce8ac5bc 100644 --- a/tortoise/__init__.py +++ b/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) diff --git a/tortoise/backends/asyncpg/client.py b/tortoise/backends/asyncpg/client.py index 08e24826c..4c2f3b9f7 100644 --- a/tortoise/backends/asyncpg/client.py +++ b/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) diff --git a/tortoise/backends/base/client.py b/tortoise/backends/base/client.py index 66ad46e69..3c9106db4 100644 --- a/tortoise/backends/base/client.py +++ b/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: """ diff --git a/tortoise/backends/base_postgres/client.py b/tortoise/backends/base_postgres/client.py index 052667098..fb83bbc2d 100644 --- a/tortoise/backends/base_postgres/client.py +++ b/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, diff --git a/tortoise/backends/psycopg/client.py b/tortoise/backends/psycopg/client.py index a9cf41601..1482bd7dc 100644 --- a/tortoise/backends/psycopg/client.py +++ b/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, }