Microsoft Entra ID (AAD)

EntraIdProvider / EntraIdPgProvider / EntraIdRedisProvider and the raw token fetchers, backed by azure.identity / azure.identity.aio. Part of the taskq[aad] extra (includes aiohttp). See the Managed Identities guide for prerequisites and a full worker example.

aad

Microsoft Entra ID (AAD) credential providers.

This module is part of the taskq[aad] optional extra. It provides :class:~taskq.auth.PgCredentialProvider and :class:~taskq.auth.RedisCredentialProvider implementations backed by Microsoft Entra ID (Azure Active Directory) managed identities, plus the raw token fetchers for users building their own providers. Install with::

pip install 'taskq-py[aad]'

Usage

::

from azure.identity.aio import DefaultAzureCredential
from taskq.auth import make_pg_pool_factory, make_redis_client_factory
from taskq.aad import EntraIdProvider

cred = DefaultAzureCredential()
provider = EntraIdProvider(cred)

WorkerConnections(
    dispatcher_pool_factory=make_pg_pool_factory(
        settings.pg_dsn_direct, provider, max_size=settings.dispatcher_pool_size,
    ),
    redis_client_factory=make_redis_client_factory(settings.redis_url, provider),
)

The provider implements both Protocols, pass the same instance to :func:~taskq.auth.make_pg_pool_factory and :func:~taskq.auth.make_redis_client_factory. For PG-only or Redis-only deployments, use :class:EntraIdPgProvider / :class:EntraIdRedisProvider individually.

Credentials

The helpers accept any object exposing get_token(*scopes) -> AccessToken , i.e. :class:azure.core.credentials.TokenCredential (sync) or its async counterpart from :mod:azure.identity.aio (e.g. :class:azure.identity.aio.DefaultAzureCredential). See :data:AadCredential. Sync credentials are offloaded to a thread so their blocking HTTP never stalls the event loop. An explicit credential is caller-owned: create it once per process (async credentials are async context managers, close them in your lifespan). Pass None to let the provider lazily create one async DefaultAzureCredential and reuse it for its lifetime (a per-fetch credential would leak unclosed aiohttp sessions and cold-cache every acquisition); that credential is the provider's own, released by await provider.aclose() (or by using the provider as an async context manager), taskq ui serve closes the providers it loads at shutdown, and an embedder closes the ones it constructs. The one-shot helpers fetch_pg_access_token / fetch_redis_credentials called without a credential create and close one per call.

This module never imports azure.identity at module top level, the import is deferred to call time so import taskq.aad is safe without the extra installed. Note import azure.identity alone does NOT make azure.identity.aio available, the subpackage must be imported explicitly.

__all__ module-attribute

__all__ = [
    "PG_TOKEN_SCOPE",
    "REDIS_TOKEN_SCOPE",
    "AadCredential",
    "EntraIdPgProvider",
    "EntraIdProvider",
    "EntraIdRedisProvider",
    "fetch_pg_access_token",
    "fetch_redis_credentials",
]

PG_TOKEN_SCOPE module-attribute

PG_TOKEN_SCOPE = (
    "https://ossrdbms-aad.database.windows.net/.default"
)

REDIS_TOKEN_SCOPE module-attribute

REDIS_TOKEN_SCOPE = 'https://redis.azure.com/.default'

AadCredential

Bases: Protocol

Structural protocol for an Azure credential with get_token.

Matches both the sync :class:azure.core.credentials.TokenCredential (get_token returns :class:~azure.core.credentials.AccessToken) and the async credentials from :mod:azure.identity.aio (get_token returns an awaitable). :func:_get_token await-detects the result so either form works.

get_token

get_token(
    *scopes: str,
    claims: str | None = None,
    tenant_id: str | None = None,
    **kwargs: Any,
) -> Any

Request an access token for scopes (sync or async).

Source code in src/taskq/aad.py

def get_token(
    self,
    *scopes: str,
    claims: str | None = None,
    tenant_id: str | None = None,
    **kwargs: Any,
) -> Any:
    """Request an access token for *scopes* (sync or async)."""
    ...

EntraIdPgProvider

EntraIdPgProvider(credential: AadCredential | None = None)

Bases: _EntraIdProviderBase, PgCredentialProvider

:class:~taskq.auth.PgCredentialProvider backed by Microsoft Entra ID.

Returns the AAD token as the Postgres password; the DSN's existing user (the AAD principal name) is preserved.

credential defaults to a lazily-created, provider-reused async DefaultAzureCredential; pass a process-wide credential to own its lifecycle (recommended in production).

Source code in src/taskq/aad.py

def __init__(self, credential: AadCredential | None = None) -> None:
    self._credential = credential
    self._default_credential = _LazyDefaultCredential()

get_pg_credential async

get_pg_credential() -> PgCredential

Source code in src/taskq/aad.py

async def get_pg_credential(self) -> PgCredential:
    token = await fetch_pg_access_token(self._resolve_credential())
    return PgCredential(password=token)

EntraIdRedisProvider

EntraIdRedisProvider(
    credential: AadCredential | None = None,
    *,
    username: str | None = None,
)

Bases: _EntraIdProviderBase, RedisCredentialProvider

:class:~taskq.auth.RedisCredentialProvider backed by Microsoft Entra ID.

Returns (managed-identity object ID, AAD token). Pass username explicitly in production to avoid JWT decoding on every reconnect.

Source code in src/taskq/aad.py

def __init__(
    self,
    credential: AadCredential | None = None,
    *,
    username: str | None = None,
) -> None:
    super().__init__(credential)
    self._username = username

get_redis_credential async

get_redis_credential() -> RedisCredential

Source code in src/taskq/aad.py

async def get_redis_credential(self) -> RedisCredential:
    username, token = await fetch_redis_credentials(
        self._resolve_credential(), username=self._username
    )
    return RedisCredential(username=username, password=token)

EntraIdProvider

EntraIdProvider(
    credential: AadCredential | None = None,
    *,
    redis_username: str | None = None,
)

Bases: EntraIdPgProvider, EntraIdRedisProvider

AAD provider implementing both PG and Redis Protocols.

Convenience class for deployments that use AAD for both Postgres and Redis, pass one instance to :func:~taskq.auth.make_pg_pool_factory and :func:~taskq.auth.make_redis_client_factory.

Source code in src/taskq/aad.py

def __init__(
    self,
    credential: AadCredential | None = None,
    *,
    redis_username: str | None = None,
) -> None:
    EntraIdPgProvider.__init__(self, credential)
    EntraIdRedisProvider.__init__(self, credential, username=redis_username)

fetch_pg_access_token async

fetch_pg_access_token(
    credential: AadCredential | None = None,
) -> str

Fetch a fresh AAD access token for Azure Database for PostgreSQL.

credential defaults to a fresh async :class:azure.identity.aio.DefaultAzureCredential that is closed before this returns (a one-shot fetch must not leak its aiohttp session); pass your own (sync or async) to reuse a process-wide credential and its token cache.

Source code in src/taskq/aad.py

async def fetch_pg_access_token(credential: AadCredential | None = None) -> str:
    """Fetch a fresh AAD access token for Azure Database for PostgreSQL.

    ``credential`` defaults to a fresh async
    :class:`azure.identity.aio.DefaultAzureCredential` that is closed
    before this returns (a one-shot fetch must not leak its aiohttp
    session); pass your own (sync or async) to reuse a process-wide
    credential and its token cache.
    """
    if credential is not None:
        return await _get_token(credential, PG_TOKEN_SCOPE)
    default = _default_credential()
    try:
        return await _get_token(default, PG_TOKEN_SCOPE)
    finally:
        await _close_credential(default)

fetch_redis_credentials async

fetch_redis_credentials(
    credential: AadCredential | None = None,
    *,
    username: str | None = None,
) -> tuple[str, str]

Fetch AAD credentials (username, password) for Azure Cache for Redis.

The password is the AAD token. The username is the managed identity's object ID, decoded from the JWT oid claim, unless username is passed explicitly (recommended in production: pass the object ID to avoid relying on JWT shape).

Source code in src/taskq/aad.py

async def fetch_redis_credentials(
    credential: AadCredential | None = None,
    *,
    username: str | None = None,
) -> tuple[str, str]:
    """Fetch AAD credentials ``(username, password)`` for Azure Cache for Redis.

    The password is the AAD token. The username is the managed identity's
    **object ID**, decoded from the JWT ``oid`` claim, unless ``username``
    is passed explicitly (recommended in production: pass the object ID to
    avoid relying on JWT shape).
    """
    if credential is not None:
        token = await _get_token(credential, REDIS_TOKEN_SCOPE)
    else:
        default = _default_credential()
        try:
            token = await _get_token(default, REDIS_TOKEN_SCOPE)
        finally:
            await _close_credential(default)
    if username is not None:
        return username, token
    oid = _decode_jwt_oid(token)
    if oid is None:
        raise ValueError(
            "Could not decode 'oid' claim from the AAD token. Pass username= "
            "explicitly with the managed identity's object ID."
        )
    return oid, token