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. The 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).

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; pass your own (sync or async) to reuse a process-wide credential.

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`; pass your own
    (sync or async) to reuse a process-wide credential.
    """
    cred = credential if credential is not None else _default_credential()
    return await _get_token(cred, PG_TOKEN_SCOPE)

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).
    """
    cred = credential if credential is not None else _default_credential()
    token = await _get_token(cred, REDIS_TOKEN_SCOPE)
    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