Skip to content

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

Open
get-bowen wants to merge 1 commit into
tortoise:developfrom
get-bowen:feat/callable-password
Open

feat(postgres): allow password to be a sync/async callable#2261
get-bowen wants to merge 1 commit into
tortoise:developfrom
get-bowen:feat/callable-password

Conversation

@get-bowen

Copy link
Copy Markdown

Description

Allow the PostgreSQL password credential to be a sync or async callable, resolved once per new connection.

Short-lived database credentials — AWS RDS/Aurora IAM tokens, Azure Entra ID tokens, Vault leases — expire long before a connection pool does. Today the only way to use them with Tortoise is to tear down and re-init the client on a timer, which kills in-flight queries when the pool is closed underneath them.

async def get_token() -> str:
    return await mint_short_lived_token()   # cached until shortly before expiry

await Tortoise.init(
    config={
        "connections": {
            "default": {
                "engine": "tortoise.backends.asyncpg",
                "credentials": {
                    "host": "db.host",
                    "port": 5432,
                    "user": "someuser",
                    "password": get_token,
                    "database": "somedb",
                },
            }
        },
        "apps": {...},
    }
)

Refs #2034.

Scope

This deliberately adds the primitive rather than native IAM support. #2034 asks for an iam_auth flag with boto3 token generation built in; that would put a cloud-vendor SDK inside Tortoise and only solve one provider. A callable password is what asyncpg and SQLAlchemy both expose, and it covers AWS IAM, Azure Entra ID, GCP, Vault and anything else in a few lines of user code.

PostgreSQL only. MySQL has the same create_pool(password=...) shape, but neither aiomysql nor asyncmy resolves callables, so supporting it there needs its own design — happy to follow up if you'd like it.

Motivation and Context

We run Tortoise against Azure Postgres with Entra ID (passwordless) auth, where tokens live ~60 minutes. Our workaround was a background task calling Tortoise.close_connections() + re-init every 5 minutes. On the asyncpg backend _close() gives in-flight work a 10-second grace period before pool.terminate(), so a long-running batch job gets its connections pulled mid-statement and fails with InterfaceError: pool is closing. Resolving the password per connection removes the need to recycle the pool at all.

How Has This Been Tested?

  • New tests/backends/test_password_factory.py: resolve_password over string/None/sync/async callables, asyncpg forwarding the callable untouched, psycopg keeping the password out of the conninfo, the generated psycopg connection class minting a fresh password per connect, and star_password not choking on a callable.
  • make check — ruff format + ruff check + mypy (134 files) + bandit all clean.
  • Full sqlite suite: 1914 passed, 148 skipped, 2 xfailed.
  • Verified end to end against a live Azure Postgres instance with a rotating Entra ID token on the asyncpg backend.

Implementation notes

  • asyncpg already resolves callable passwords itself (connect_utils.py, awaits the result if awaitable), and BasePostgresClient forwarded password untouched — so this worked by accident but was typed str | None and undocumented. It is now typed, tested and documented so it doesn't get "fixed" away.
  • psycopg bakes credentials into an immutable conninfo string, so the password is instead injected at connect time through a generated connection class. psycopg_pool does support a callable conninfo, but only from 3.3, while pyproject.toml allows psycopg[pool] >=3.0.12 — this keeps that floor intact.
  • Tortoise.star_password() was a latent crash: it does password[0 : len(password) // 3] and str.replace(password, ...), both TypeError on a non-string. Now skips non-strings.

Checklist:

  • My code follows the code style of this project.
  • My change requires a change to the documentation.
  • I have updated the documentation accordingly.
  • I have added tests to cover my changes.
  • All new and existing tests passed.
  • I have added an entry to the CHANGELOG.

Short-lived database credentials (AWS RDS/Aurora IAM tokens, Azure Entra ID
tokens, Vault leases) expire long before a connection pool does. Allow the
postgres 'password' credential to be a callable, resolved once per new
connection, so tokens are refreshed transparently.

- asyncpg resolves the callable natively; it is now forwarded with the right type
- psycopg bakes credentials into an immutable conninfo string, so the password is
  injected at connect time via a generated connection class instead
- Tortoise.star_password no longer crashes on a non-string password

Refs tortoise#2034
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant