feat(db): add Azure managed identity authentication for PostgreSQL

[RogerHYang]

Resolves #9069. Supersedes #12667 (will be closed after this lands).

Summary

Adds first-class support for Microsoft Entra managed-identity authentication when connecting Phoenix to Azure Database for PostgreSQL (Flexible Server). While wiring that in, the existing AWS RDS IAM auth path is modernized so both clouds share a consistent shape — the on-disk module is renamed to aws_auth.py, token generation is switched from blocking boto3 to aioboto3, and the misconceived "align pool_recycle to the token lifetime" coupling is replaced with a plain hygiene cap on every engine.

What's new

Azure managed identity (feature)

• New module src/phoenix/db/azure_auth.py: create_azure_token_connection_creator() returns an async creator that obtains a fresh JWT via credential.get_token(scope) on every new DB connection. Token caching and refresh are handled by azure-identity's built-in MSAL TokenCache (sync and async share the same machinery via ManagedIdentityClientBase, verified by reading the installed 1.18.0 and 1.25.3 source).
• New env vars:
  • PHOENIX_POSTGRES_USE_AZURE_MANAGED_IDENTITY (bool, default false).
  • PHOENIX_POSTGRES_AZURE_SCOPE (string, default https://ossrdbms-aad.database.windows.net/.default). Only needed for sovereign clouds, e.g. Azure US Government.
• Validation is enforced at the layer that actually uses each input: mutual exclusion with AWS IAM (cannot enable both simultaneously) at the engines.py dispatch point; PASSWORD-must-not-be-set and USER-must-be-set checks in get_env_postgres_connection_str alongside URL composition; the azure-identity import check in azure_auth.py next to the code that uses it.
• New optional [azure] extra pinning azure-identity >= 1.18.0 and aiohttp. 1.18.0 is the first stable release where ManagedIdentityCredential is implemented on top of MSAL's TokenCache (source-verified).

AWS RDS IAM (modernized)

• Rename src/phoenix/db/iam_auth.py → src/phoenix/db/aws_auth.py so the module name reflects the cloud rather than the auth mechanism. Tracked as a rename by git.
• Switch token generator from blocking boto3 to aioboto3 so generate_db_auth_token no longer blocks the event loop during connection-open. One aioboto3.Session is reused for the factory's lifetime; a fresh client context is opened per connection because aioboto3 does not allow reusing the same async-context object.
• New public factory create_aws_iam_token_connection_creator() mirroring the Azure one (validation + closed-over async creator). _generate_aws_rds_token is now private.

Pool configuration

• _POOL_RECYCLE_SECONDS = 3300 (module-level constant in engines.py) applied consistently as a hygiene cap across the Azure, AWS, and plain-password main engines. This is not a token-expiry mechanism — PostgreSQL validates the credential only during the startup handshake and never re-checks it mid-session, and both Microsoft and AWS explicitly document that existing connections are unaffected by later token expiry. See internal_docs/specs/postgres-cloud-auth-pooling.md for the full analysis with primary-source citations.
• pool_pre_ping=True on every main engine to catch stale connections (idle-timeout drops, failovers, TCP resets) at checkout rather than surfacing them as errors to callers.
• NullPool migration engines intentionally omit both knobs — SQLAlchemy skips the pre-ping on freshly-created connections (verified in sqlalchemy/pool/base.py), and NullPool creates a fresh record per checkout so pool_recycle has nothing to age out.
• Cloud-auth imports are hoisted inside their respective branches so importing engines.py does not require azure-identity or aioboto3 unless the corresponding feature is actually used.
• Dispatch-time mutual exclusion check raises ValueError when both PHOENIX_POSTGRES_USE_AWS_IAM_AUTH and PHOENIX_POSTGRES_USE_AZURE_MANAGED_IDENTITY are set to true, so the misconfiguration fails loudly instead of silently picking one branch.

Tests

• tests/unit/db/test_azure_auth.py — 6 parsimonious tests covering validation, connect_args forwarding, per-call token retrieval with password wiring, and scope-from-env integration.
• tests/unit/db/test_aws_auth.py — new 6-test suite mirroring the Azure structure, including an AWS-specific test that verifies URL defaults (port=5432, database='postgres') kick in when the URL omits them.
• tests/unit/test_config.py — new TestGetEnvPostgresAzureScope + TestGetEnvPostgresUseManagedIdentity; MI-specific validation is exercised via the existing TestPostgresConnectionStringManagedIdentity class (mutual exclusion, password-conflict, host/user handling), which covers the layer that actually composes the connection string.

Documentation

• New spec internal_docs/specs/postgres-cloud-auth-pooling.md: a 384-line reference covering PostgreSQL protocol semantics, Azure managed-identity flow, AWS RDS IAM flow, a side-by-side comparison of how the two clouds actually differ (bearer tokens vs HMAC signing keys; both have a metadata-service call at credential bootstrap), SQLAlchemy pool mechanics, common assumptions that don't hold, and the full security model. Every non-trivial claim cites a primary source; inferences are explicitly marked.
• New user guide docs/phoenix/self-hosting/configuration/using-azure-managed-identity.mdx with env-var setup, a note on non-public-cloud scope overrides, and an explicit note that no token-lifetime tuning variable is needed.
• docs/phoenix/self-hosting/configuration/using-amazon-aurora.mdx updated to remove the stale PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS example.
• docs.json + docs/phoenix/llms.txt wired for the new Azure guide.

Helm chart

• database.postgres.awsIamTokenLifetimeSeconds removed from values.yaml, the README table, and the configmap template. Custom values files that still carry the old key will have helm upgrade succeed without error; the key simply becomes unreferenced.
• scripts/ci/test_helm.py drops the custom-token-lifetime test case and the token_lifetime parameter from DatabaseValidators.aws_iam_auth.

Deprecations (backward compatible — no user action required)

• PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS is silently ignored, with a one-line startup warning so users notice. It previously controlled pool_recycle on the IAM branch; the new code uses a hygiene cap instead. Users who had pinned this to a short value (e.g. 840s) will see fewer connection rotations per hour — warmer pool, fewer TCP+TLS+PG startup handshakes, fewer SigV4 signing calls.
• helm database.postgres.awsIamTokenLifetimeSeconds is removed. helm upgrade still succeeds if a custom values file still sets it.

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated (UTC)
arize-phoenix	🟢 Ready	View Preview	Apr 15, 2026, 4:03 PM

💡 Tip: Enable Workflows to automatically generate PRs for you.

[claude]

The early return here means that when PHOENIX_POSTGRES_USE_AWS_IAM_AUTH=true (and Azure MI is not enabled), this function returns immediately with no AWS-specific startup validation at all.

The old _validate_iam_auth_config() validated three things at startup for the AWS path:

1. That boto3 was importable (clear error message if missing)
2. That PHOENIX_POSTGRES_USER was set
3. That AWS credentials were valid via sts.get_caller_identity()

After this refactor, none of those checks remain for the AWS-only path. A misconfigured AWS setup (missing aioboto3, missing user, invalid credentials) will now surface as a cryptic connection error on first DB access rather than a clear startup failure.

Consider adding an AWS-specific validation branch analogous to the Azure one:

    if not get_env_postgres_use_azure_managed_identity():
        if get_env_postgres_use_iam_auth():
            _validate_aws_iam_config()  # new helper with aioboto3 import check + user check
        return

Or, at minimum, reuse the existing user check and add an aioboto3 import check for the AWS path before returning.

See old _validate_iam_auth_config in this diff for the checks that were removed.

[RogerHYang]

Intentional. The three checks the old _validate_iam_auth_config() did are now either redundant or low-value:

1. aioboto3 importable — aws_auth.py does this at factory construction with an explicit install-instructions ImportError. The factory runs during app startup, before migrations, so re-checking it in config.py would duplicate the same message a few frames earlier.

2. PHOENIX_POSTGRES_USER set — aws_auth.py raises ValueError("Database user is required for AWS RDS IAM authentication") at factory construction. Same check, same message, right next to the code that uses the value.

3. sts.get_caller_identity() credential pre-flight — dropped deliberately. aioboto3.generate_db_auth_token surfaces credential failures (NoCredentialsError, ClientError) at the first DB connection, which is the migration step running moments after startup. The startup STS call also adds a synchronous network round-trip and an STS dependency that isn't otherwise on the critical path.

Net: the checks that matter already live at the factory layer; the one that's genuinely gone was providing a slightly-earlier version of the same information the SDK already emits.

[claude]

Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

[mitch-crites-pm]

@RogerHYang I hit a production issue testing the original PR that may affect this one as well.
DefaultAzureCredential() is created at factory call time (before the asyncio event loop is running), which causes
"RuntimeError: Event loop is closed" when the pool opens new connections. I'll pull this branch down and deploy overnight to confirm. Will report back.

[RogerHYang]

@mitch-crites-pm Thanks for the heads-up. I'll take a look at that issue.

[RogerHYang]

Root Cause Analysis

Event-loop affinity of azure.identity.aio credentials

azure.identity.aio.DefaultAzureCredential (and the inner ManagedIdentityCredential / ImdsCredential it resolves to in Azure-hosted environments) is an async-transport credential: on first get_token(scope) it lazily constructs an AsyncPipeline built around an azure.core.pipeline.transport.AioHttpTransport, which in turn opens an aiohttp.ClientSession. aiohttp.ClientSession is bound to whichever event loop was running when it was created — reusing it from a different loop (including a loop that has since been closed) raises RuntimeError: Event loop is closed at the first await inside the session.

Constructing DefaultAzureCredential() itself does not bind to a loop. The binding happens at the first await credential.get_token(...), when the pipeline is instantiated. So a credential built in sync context (before any loop exists) is fine — until some loop calls get_token on it and thereby pins the credential's transport to that loop.

This matters for any application that runs one short-lived asyncio.run(...) (e.g. a migration runner) and then later runs a long-lived loop (e.g. a web server) while sharing the same credential instance between them. The sequence is:

1. Main thread/process builds one DefaultAzureCredential in sync context. No loop is bound yet.
2. Short-lived loop A starts (typically via asyncio.run inside a worker thread for one-shot work), calls credential.get_token(scope), and the credential's aiohttp session is lazily created and bound to loop A.
3. Loop A finishes its work and is closed by asyncio.run's cleanup.
4. Long-lived loop B starts. Application code calls credential.get_token(scope) on the same credential. The credential reuses its cached aiohttp session, which is still bound to the (now closed) loop A. The next network read/write inside the session raises RuntimeError: Event loop is closed.

Symptom: RuntimeError: Event loop is closed on pool connection-open in Phoenix

Phoenix's primary DB engine and migration DB engine are both built from src/phoenix/db/engines.py's create_engine, which on the use_azure branch calls create_azure_token_connection_creator(...) once and passes the returned async_creator into both create_async_engine calls (the primary engine and the migration engine). The single closure holds one DefaultAzureCredential instance.

The migration engine runs via migrate_in_thread → asyncio.run(run()) in a worker thread (src/phoenix/db/migrate.py), which creates and then closes a fresh loop for the duration of the migration. When the first migration checkout fires async_creator(), the credential's aiohttp session is bound to that throwaway loop. When migrations finish, asyncio.run closes the loop and returns, but the credential — held alive by the closure, which is held alive by the primary engine — persists.

Later, uvicorn starts the long-lived server loop and the primary pool opens its first real connection. async_creator() calls credential.get_token(scope), azure-identity routes the call through the cached pipeline, the pipeline hands it to the aiohttp session bound to the dead migration loop, and the call raises RuntimeError: Event loop is closed. Because this surfaces inside SQLAlchemy's pool-open path, it looks like a pool error rather than a credential error.

The TokenCache-level cache cannot short-circuit the failure because the cache lookup still traverses the async pipeline on every get_token call (see the paragraph above).

[github-actions]

Card links check

No broken Card links found. Checked external links in 14.5s

[RogerHYang]

@mitch-crites-pm I fixed the event loop issue on this branch.