azure-postgresql-auth 1.0.0__py3-none-any.whl

azure_postgresql_auth/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Azure PostgreSQL Entra ID Integration Library
+
+This library provides connection classes for using Azure Entra ID authentication
+with Azure Database for PostgreSQL across different PostgreSQL drivers.
+
+Available modules (with optional dependencies):
+- psycopg2: Support for psycopg2 driver (install with: pip install azurepg-entra[psycopg2])
+- psycopg3: Support for psycopg (v3) driver (install with: pip install azurepg-entra[psycopg3])
+- sqlalchemy: Support for SQLAlchemy ORM (install with: pip install azurepg-entra[sqlalchemy])
+
+Core dependencies (always available):
+- azure-identity: For Azure Entra ID authentication
+- azure-core: Core Azure SDK functionality
+"""
+
+__version__ = "0.1.0"
+__author__ = "Microsoft Corporation"

azure_postgresql_auth/core.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+import json
+from typing import Any, cast
+
+from azure.core.credentials import TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity import CredentialUnavailableError
+from azure.identity import DefaultAzureCredential as DefaultAzureCredential
+from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
+
+from azure_postgresql_auth.errors import (
+    ScopePermissionError,
+    TokenDecodeError,
+    UsernameExtractionError,
+)
+
+AZURE_DB_FOR_POSTGRES_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
+AZURE_MANAGEMENT_SCOPE = "https://management.azure.com/.default"
+
+
+def get_entra_token(credential: TokenCredential | None, scope: str) -> str:
+    """Acquires an Entra authentication token for Azure PostgreSQL synchronously.
+
+    Parameters:
+        credential (TokenCredential or None): Credential object used to obtain the token.
+            If None, the default Azure credentials are used.
+        scope (str): The scope for the token request.
+
+    Returns:
+        str: The acquired authentication token to be used as the database password.
+    """
+    credential = credential or DefaultAzureCredential()
+    cred = credential.get_token(scope)
+    return cred.token
+
+
+async def get_entra_token_async(
+    credential: AsyncTokenCredential | None, scope: str
+) -> str:
+    """Asynchronously acquires an Entra authentication token for Azure PostgreSQL.
+
+    Parameters:
+        credential (AsyncTokenCredential or None): Asynchronous credential used to obtain the token.
+            If None, the default Azure credentials are used.
+        scope (str): The scope for the token request.
+
+    Returns:
+        str: The acquired authentication token to be used as the database password.
+    """
+    credential = credential or AsyncDefaultAzureCredential()
+    async with credential:
+        cred = await credential.get_token(scope)
+        return cred.token
+
+
+def decode_jwt(token: str) -> dict[str, Any]:
+    """Decodes a JWT token to extract its payload claims.
+
+    Parameters:
+        token (str): The JWT token string in the standard three-part format.
+
+    Returns:
+        dict[str, Any]: A dictionary containing the claims extracted from the token payload.
+
+    Raises:
+        TokenValueError: If the token format is invalid or cannot be decoded.
+    """
+    try:
+        payload = token.split(".")[1]
+        padding = "=" * (4 - len(payload) % 4)
+        decoded_payload = base64.urlsafe_b64decode(payload + padding)
+        return cast(dict[str, Any], json.loads(decoded_payload))
+    except Exception as e:
+        raise TokenDecodeError("Invalid JWT token format") from e
+
+
+def parse_principal_name(xms_mirid: str) -> str | None:
+    """Parses the principal name from an Azure resource path.
+
+    Parameters:
+        xms_mirid (str): The xms_mirid claim value containing the Azure resource path.
+
+    Returns:
+        str | None: The extracted principal name, or None if parsing fails.
+    """
+    if not xms_mirid:
+        return None
+
+    # Parse the xms_mirid claim which looks like
+    # /subscriptions/{subId}/resourcegroups/{resourceGroup}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{principalName}
+    last_slash_index = xms_mirid.rfind("/")
+    if last_slash_index == -1:
+        return None
+
+    beginning = xms_mirid[:last_slash_index]
+    principal_name = xms_mirid[last_slash_index + 1 :]
+
+    if not principal_name or not beginning.lower().endswith(
+        "providers/microsoft.managedidentity/userassignedidentities"
+    ):
+        return None
+
+    return principal_name
+
+
+def get_entra_conninfo(credential: TokenCredential | None) -> dict[str, str]:
+    """Synchronously obtains connection information from Entra authentication for Azure PostgreSQL.
+
+    This function acquires an access token from Azure Entra ID and extracts the username
+    from the token claims. It tries multiple claim sources to determine the username.
+
+    Parameters:
+        credential (TokenCredential or None): The credential used for token acquisition.
+            If None, DefaultAzureCredential() is used to automatically discover credentials.
+
+    Returns:
+        dict[str, str]: A dictionary with 'user' and 'password' keys, where:
+            - 'user': The extracted username from token claims
+            - 'password': The Entra ID access token for database authentication
+
+    Raises:
+        TokenDecodeError: If the JWT token cannot be decoded or is malformed.
+        UsernameExtractionError: If the username cannot be extracted from token claims.
+        ScopePermissionError: The token could not be acquired from the management scope, possibly due to insufficient permissions.
+    """
+    credential = credential or DefaultAzureCredential()
+
+    # Always get the DB-scope token for password
+    db_token = get_entra_token(credential, AZURE_DB_FOR_POSTGRES_SCOPE)
+    try:
+        db_claims = decode_jwt(db_token)
+    except TokenDecodeError:
+        raise
+    xms_mirid = db_claims.get("xms_mirid")
+    username = (
+        parse_principal_name(xms_mirid)
+        if isinstance(xms_mirid, str)
+        else None
+        or db_claims.get("upn")
+        or db_claims.get("preferred_username")
+        or db_claims.get("unique_name")
+    )
+
+    if not username:
+        # Fall back to management scope ONLY to discover username
+        try:
+            mgmt_token = get_entra_token(credential, AZURE_MANAGEMENT_SCOPE)
+        except (CredentialUnavailableError, ClientAuthenticationError) as e:
+            raise ScopePermissionError(
+                "Failed to acquire token from management scope"
+            ) from e
+        try:
+            mgmt_claims = decode_jwt(mgmt_token)
+        except TokenDecodeError:
+            raise
+        xms_mirid = mgmt_claims.get("xms_mirid")
+        username = (
+            parse_principal_name(xms_mirid)
+            if isinstance(xms_mirid, str)
+            else None
+            or mgmt_claims.get("upn")
+            or mgmt_claims.get("preferred_username")
+            or mgmt_claims.get("unique_name")
+        )
+
+    if not username:
+        raise UsernameExtractionError(
+            "Could not determine username from token claims. "
+            "Ensure the identity has the proper Azure AD attributes."
+        )
+
+    return {"user": username, "password": db_token}
+
+
+async def get_entra_conninfo_async(
+    credential: AsyncTokenCredential | None,
+) -> dict[str, str]:
+    """Asynchronously obtains connection information from Entra authentication for Azure PostgreSQL.
+
+    This function acquires an access token from Azure Entra ID and extracts the username
+    from the token claims. It tries multiple claim sources to determine the username.
+
+    Parameters:
+        credential (AsyncTokenCredential or None): The async credential used for token acquisition.
+            If None, AsyncDefaultAzureCredential() is used to automatically discover credentials.
+
+    Returns:
+        dict[str, str]: A dictionary with 'user' and 'password' keys, where:
+            - 'user': The extracted username from token claims
+            - 'password': The Entra ID access token for database authentication
+
+    Raises:
+        TokenDecodeError: If the JWT token cannot be decoded or is malformed.
+        UsernameExtractionError: If the username cannot be extracted from token claims.
+        ScopePermissionError: The token could not be acquired from the management scope, possibly due to insufficient permissions.
+    """
+    credential = credential or AsyncDefaultAzureCredential()
+
+    db_token = await get_entra_token_async(credential, AZURE_DB_FOR_POSTGRES_SCOPE)
+    try:
+        db_claims = decode_jwt(db_token)
+    except TokenDecodeError:
+        raise
+    xms_mirid = db_claims.get("xms_mirid")
+    username = (
+        parse_principal_name(xms_mirid)
+        if isinstance(xms_mirid, str)
+        else None
+        or db_claims.get("upn")
+        or db_claims.get("preferred_username")
+        or db_claims.get("unique_name")
+    )
+
+    if not username:
+        try:
+            mgmt_token = await get_entra_token_async(credential, AZURE_MANAGEMENT_SCOPE)
+        except (CredentialUnavailableError, ClientAuthenticationError) as e:
+            raise ScopePermissionError(
+                "Failed to acquire token from management scope"
+            ) from e
+        try:
+            mgmt_claims = decode_jwt(mgmt_token)
+        except TokenDecodeError:
+            raise
+        xms_mirid = mgmt_claims.get("xms_mirid")
+        username = (
+            parse_principal_name(xms_mirid)
+            if isinstance(xms_mirid, str)
+            else None
+            or mgmt_claims.get("upn")
+            or mgmt_claims.get("preferred_username")
+            or mgmt_claims.get("unique_name")
+        )
+
+    if not username:
+        raise UsernameExtractionError(
+            "Could not determine username from token claims. "
+            "Ensure the identity has the proper Azure AD attributes."
+        )
+
+    return {"user": username, "password": db_token}

azure_postgresql_auth/errors.py

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+class AzurePgEntraError(Exception):
+    """Base class for all custom exceptions in the project."""
+
+    pass
+
+
+class TokenDecodeError(AzurePgEntraError):
+    """Raised when a token value is invalid."""
+
+    pass
+
+
+class UsernameExtractionError(AzurePgEntraError):
+    """Raised when username cannot be extracted from token."""
+
+    pass
+
+
+class CredentialValueError(AzurePgEntraError):
+    """Raised when token credential is invalid."""
+
+    pass
+
+
+class EntraConnectionValueError(AzurePgEntraError):
+    """Raised when Entra connection credentials are invalid."""
+
+    pass
+
+
+class ScopePermissionError(AzurePgEntraError):
+    """Raised when the provided scope does not have sufficient permissions."""
+
+    pass

azure_postgresql_auth/psycopg2/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Psycopg2 support for Azure Entra ID authentication with Azure Database for PostgreSQL.
+
+This module provides a connection class that handles Azure Entra ID token acquisition
+and authentication for synchronous PostgreSQL connections.
+
+Requirements:
+    Install with: pip install azurepg-entra[psycopg2]
+
+    This will install:
+    - psycopg2-binary>=2.9.0
+
+Classes:
+    EntraConnection: Synchronous connection class with Entra ID authentication
+"""
+
+from .entra_connection import (
+    EntraConnection,
+)
+
+__all__ = [
+    "EntraConnection",
+]

azure_postgresql_auth/psycopg2/entra_connection.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from azure.core.credentials import TokenCredential
+
+from azure_postgresql_auth.core import get_entra_conninfo
+from azure_postgresql_auth.errors import (
+    CredentialValueError,
+    EntraConnectionValueError,
+)
+
+try:
+    from psycopg2.extensions import connection, make_dsn, parse_dsn
+except ImportError as e:
+    # Provide a helpful error message if psycopg2 dependencies are missing
+    raise ImportError(
+        "psycopg2 dependencies are not installed. "
+        "Install them with: pip install azurepg-entra[psycopg2]"
+    ) from e
+
+class EntraConnection(connection):
+    """Establishes a synchronous PostgreSQL connection using Entra authentication.
+
+    This connection class automatically acquires Azure Entra ID credentials when user
+    or password are not provided in the DSN or connection parameters. Authentication
+    errors are printed to console for debugging purposes.
+
+    Parameters:
+        dsn (str): PostgreSQL connection string (Data Source Name).
+        **kwargs: Additional keyword arguments including:
+            - credential (TokenCredential, optional): Azure credential for token acquisition.
+              If None, DefaultAzureCredential() is used.
+            - user (str, optional): Database username. If not provided, extracted from Entra token.
+            - password (str, optional): Database password. If not provided, uses Entra access token.
+
+    Raises:
+        CredentialValueError: If the provided credential is not a valid TokenCredential.
+        EntraConnectionValueError: If Entra connection credentials cannot be retrieved
+    """
+
+    def __init__(self, dsn: str, **kwargs: Any) -> None:
+        # Extract current DSN params
+        dsn_params = parse_dsn(dsn) if dsn else {}
+
+        credential = kwargs.pop("credential", None)
+        if credential and not isinstance(credential, (TokenCredential)):
+            raise CredentialValueError(
+                "credential must be a TokenCredential for sync connections"
+            )
+
+        # Check if user and password are already provided
+        has_user = "user" in dsn_params or "user" in kwargs
+        has_password = "password" in dsn_params or "password" in kwargs
+
+        # Only get Entra credentials if user or password is missing
+        if not has_user or not has_password:
+            try:
+                entra_creds = get_entra_conninfo(credential)
+            except (Exception) as e:
+                raise EntraConnectionValueError(
+                    "Could not retrieve Entra credentials"
+                ) from e
+
+            # Only update missing credentials
+            if not has_user and "user" in entra_creds:
+                dsn_params["user"] = entra_creds["user"]
+            if not has_password and "password" in entra_creds:
+                dsn_params["password"] = entra_creds["password"]
+
+        # Update DSN params with any kwargs (kwargs take precedence)
+        dsn_params.update(kwargs)
+
+        # Create new DSN with updated credentials
+        new_dsn = make_dsn(**dsn_params)
+
+        # Call parent constructor with updated DSN only
+        super().__init__(new_dsn)

azure_postgresql_auth/psycopg3/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Psycopg3 support for Azure Entra ID authentication with Azure Database for PostgreSQL.
+
+This module provides connection classes that extend psycopg's Connection and AsyncConnection
+to automatically handle Azure Entra ID token acquisition and authentication.
+
+Requirements:
+    Install with: pip install azurepg-entra[psycopg3]
+
+    This will install:
+    - psycopg[binary]>=3.1.0
+    - aiohttp>=3.8.0
+
+Classes:
+    EntraConnection: Synchronous connection class with Entra ID authentication
+    AsyncEntraConnection: Asynchronous connection class with Entra ID authentication
+"""
+
+from .async_entra_connection import AsyncEntraConnection
+from .entra_connection import EntraConnection
+
+__all__ = ["EntraConnection", "AsyncEntraConnection"]

azure_postgresql_auth/psycopg3/async_entra_connection.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from azure.core.credentials_async import AsyncTokenCredential
+
+try:
+    from psycopg import AsyncConnection
+except ImportError as e:
+    raise ImportError(
+        "psycopg3 dependencies are not installed. "
+        "Install them with: pip install azurepg-entra[psycopg3]"
+    ) from e
+
+from azure_postgresql_auth.core import get_entra_conninfo_async
+from azure_postgresql_auth.errors import (
+    CredentialValueError,
+    EntraConnectionValueError,
+)
+
+
+class AsyncEntraConnection(AsyncConnection):
+    """Asynchronous connection class for using Entra authentication with Azure PostgreSQL."""
+
+    @classmethod
+    async def connect(cls, *args: Any, **kwargs: Any) -> "AsyncEntraConnection":
+        """Establishes an asynchronous PostgreSQL connection using Entra authentication.
+
+        This method automatically acquires Azure Entra ID credentials when user or password
+        are not provided in the connection parameters. Authentication errors are printed to
+        console for debugging purposes.
+
+        Parameters:
+            *args: Positional arguments to be forwarded to the parent connection method.
+            **kwargs: Keyword arguments including:
+                - credential (AsyncTokenCredential, optional): Async Azure credential for token acquisition.
+                - user (str, optional): Database username. If not provided, extracted from Entra token.
+                - password (str, optional): Database password. If not provided, uses Entra access token.
+
+        Returns:
+            AsyncEntraConnection: An open asynchronous connection to the PostgreSQL database.
+
+        Raises:
+            CredentialValueError: If the provided credential is not a valid AsyncTokenCredential.
+            EntraConnectionValueError: If Entra connection credentials are invalid.
+        """
+        credential = kwargs.pop("credential", None)
+        if credential and not isinstance(credential, (AsyncTokenCredential)):
+            raise CredentialValueError(
+                "credential must be an AsyncTokenCredential for async connections"
+            )
+
+        # Check if we need to acquire Entra authentication info
+        if not kwargs.get("user") or not kwargs.get("password"):
+            try:
+                entra_conninfo = await get_entra_conninfo_async(credential)
+            except Exception as e:
+                raise EntraConnectionValueError(
+                    "Could not retrieve Entra credentials"
+                ) from e
+            # Always use the token password when Entra authentication is needed
+            kwargs["password"] = entra_conninfo["password"]
+            if not kwargs.get("user"):
+                # If user isn't already set, use the username from the token
+                kwargs["user"] = entra_conninfo["user"]
+        return await super().connect(*args, **kwargs)

azure_postgresql_auth/psycopg3/entra_connection.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from azure.core.credentials import TokenCredential
+
+try:
+    from psycopg import Connection
+except ImportError as e:
+    raise ImportError(
+        "psycopg3 dependencies are not installed. "
+        "Install them with: pip install azurepg-entra[psycopg3]"
+    ) from e
+
+from azure_postgresql_auth.core import get_entra_conninfo
+from azure_postgresql_auth.errors import (
+    CredentialValueError,
+    EntraConnectionValueError,
+)
+
+
+class EntraConnection(Connection):
+    """Synchronous connection class for using Entra authentication with Azure PostgreSQL."""
+
+    @classmethod
+    def connect(cls, *args: Any, **kwargs: Any) -> "EntraConnection":
+        """Establishes a synchronous PostgreSQL connection using Entra authentication.
+
+        This method automatically acquires Azure Entra ID credentials when user or password
+        are not provided in the connection parameters. If authentication fails, the original
+        exception is re-raised to the caller.
+
+        Parameters:
+            *args: Positional arguments to be forwarded to the parent connection method.
+            **kwargs: Keyword arguments including:
+                - credential (TokenCredential, optional): Azure credential for token acquisition.
+                - user (str, optional): Database username. If not provided, extracted from Entra token.
+                - password (str, optional): Database password. If not provided, uses Entra access token.
+
+        Returns:
+            EntraConnection: An open synchronous connection to the PostgreSQL database.
+
+        Raises:
+            CredentialValueError: If the provided credential is not a valid TokenCredential.
+            EntraConnectionValueError: If Entra connection credentials cannot be retrieved
+        """
+        credential = kwargs.pop("credential", None)
+        if credential and not isinstance(credential, (TokenCredential)):
+            raise CredentialValueError(
+                "credential must be a TokenCredential for sync connections"
+            )
+
+        # Check if we need to acquire Entra authentication info
+        if not kwargs.get("user") or not kwargs.get("password"):
+            try:
+                entra_conninfo = get_entra_conninfo(credential)
+            except Exception as e:
+                raise EntraConnectionValueError(
+                    "Could not retrieve Entra credentials"
+                ) from e
+            # Always use the token password when Entra authentication is needed
+            kwargs["password"] = entra_conninfo["password"]
+            if not kwargs.get("user"):
+                # If user isn't already set, use the username from the token
+                kwargs["user"] = entra_conninfo["user"]
+        return super().connect(*args, **kwargs)

azure_postgresql_auth/sqlalchemy/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+SQLAlchemy integration for Azure PostgreSQL with Entra ID authentication.
+
+This module provides integration between SQLAlchemy and Azure Entra ID
+authentication for PostgreSQL connections. It automatically handles token acquisition
+and credential injection through SQLAlchemy's event system.
+
+Requirements:
+    Install with: pip install azurepg-entra[sqlalchemy]
+
+    This will install:
+    - sqlalchemy>=2.0.0
+    - aiohttp>=3.8.0
+
+Functions:
+    enable_entra_authentication: Enable Entra ID authentication for synchronous SQLAlchemy engines
+    enable_entra_authentication_async: Enable Entra ID authentication for asynchronous SQLAlchemy engines
+"""
+
+from .async_entra_connection import enable_entra_authentication_async
+from .entra_connection import enable_entra_authentication
+
+__all__ = [
+    "enable_entra_authentication",
+    "enable_entra_authentication_async",
+]

azure_postgresql_auth/sqlalchemy/async_entra_connection.py

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from azure.core.credentials import TokenCredential
+
+try:
+    from sqlalchemy import event
+    from sqlalchemy.engine import Dialect
+    from sqlalchemy.ext.asyncio import AsyncEngine
+except ImportError as e:
+    # Provide a helpful error message if SQLAlchemy dependencies are missing
+    raise ImportError(
+        "SQLAlchemy dependencies are not installed. "
+        "Install them with: pip install azurepg-entra[sqlalchemy]"
+    ) from e
+
+from azure_postgresql_auth.core import get_entra_conninfo
+from azure_postgresql_auth.errors import (
+    CredentialValueError,
+    EntraConnectionValueError,
+)
+
+
+def enable_entra_authentication_async(engine: AsyncEngine) -> None:
+    """
+    Enable Azure Entra ID authentication for an async SQLAlchemy engine.
+
+    This function registers an event listener that automatically provides
+    Entra ID credentials for each database connection if they are not already set.
+    Event handlers do not support async behavior so the token fetching will still
+    be synchronous.
+
+    Args:
+        engine: The async SQLAlchemy Engine to enable Entra authentication for
+    """
+
+    @event.listens_for(engine.sync_engine, "do_connect")
+    def provide_token(
+        dialect: Dialect, conn_rec: Any, cargs: Any, cparams: dict[str, Any]
+    ) -> None:
+        """Event handler that provides Entra credentials for each sync connection.
+
+        Raises:
+            CredentialValueError: If the provided credential is not a valid TokenCredential.
+            EntraConnectionValueError: If Entra connection credentials cannot be retrieved
+        """
+        credential = cparams.get("credential", None)
+        if credential and not isinstance(credential, (TokenCredential)):
+            raise CredentialValueError(
+                "credential must be a TokenCredential for async connections"
+            )
+        # Check if credentials are already present
+        has_user = "user" in cparams
+        has_password = "password" in cparams
+
+        # Only get Entra credentials if user or password is missing
+        if not has_user or not has_password:
+            try:
+                entra_creds = get_entra_conninfo(credential)
+            except Exception as e:
+                raise EntraConnectionValueError(
+                    "Could not retrieve Entra credentials"
+                ) from e
+            # Only update missing credentials
+            if not has_user and "user" in entra_creds:
+                cparams["user"] = entra_creds["user"]
+            if not has_password and "password" in entra_creds:
+                cparams["password"] = entra_creds["password"]
+
+        # Strip helper-only param before DBAPI connect to avoid 'invalid connection option'
+        if "credential" in cparams:
+            del cparams["credential"]

azure_postgresql_auth/sqlalchemy/entra_connection.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from azure.core.credentials import TokenCredential
+
+try:
+    from sqlalchemy import Engine, event
+    from sqlalchemy.engine import Dialect
+except ImportError as e:
+    raise ImportError(
+        "SQLAlchemy dependencies are not installed. "
+        "Install them with: pip install azurepg-entra[sqlalchemy]"
+    ) from e
+
+from azure_postgresql_auth.core import get_entra_conninfo
+from azure_postgresql_auth.errors import (
+    CredentialValueError,
+    EntraConnectionValueError,
+)
+
+
+def enable_entra_authentication(engine: Engine) -> None:
+    """
+    Enable Azure Entra ID authentication for a SQLAlchemy engine.
+
+    This function registers an event listener that automatically provides
+    Entra ID credentials for each database connection if they are not already set.
+
+    Args:
+        engine: The SQLAlchemy Engine to enable Entra authentication for
+    """
+
+    @event.listens_for(engine, "do_connect")
+    def provide_token(
+        dialect: Dialect, conn_rec: Any, cargs: Any, cparams: dict[str, Any]
+    ) -> None:
+        """Event handler that provides Entra credentials for each connection.
+
+        Raises:
+            CredentialValueError: If the provided credential is not a valid TokenCredential.
+            EntraConnectionValueError: If Entra connection credentials cannot be retrieved
+        """
+        credential = cparams.get("credential", None)
+        if credential and not isinstance(credential, (TokenCredential)):
+            raise CredentialValueError(
+                "credential must be a TokenCredential for sync connections"
+            )
+        # Check if credentials are already present
+        has_user = "user" in cparams
+        has_password = "password" in cparams
+
+        # Only get Entra credentials if user or password is missing
+        if not has_user or not has_password:
+            try:
+                entra_creds = get_entra_conninfo(credential)
+            except Exception as e:
+                raise EntraConnectionValueError(
+                    "Could not retrieve Entra credentials"
+                ) from e
+            # Only update missing credentials
+            if not has_user and "user" in entra_creds:
+                cparams["user"] = entra_creds["user"]
+            if not has_password and "password" in entra_creds:
+                cparams["password"] = entra_creds["password"]
+
+        # Remove the helper-only parameter so the DBAPI (psycopg/psycopg2) doesn't see an
+        # unknown connection option and raise 'invalid connection option "credential"'.
+        if "credential" in cparams:
+            del cparams["credential"]

azure_postgresql_auth-1.0.0.dist-info/METADATA

@@ -0,0 +1,360 @@
+Metadata-Version: 2.4
+Name: azure-postgresql-auth
+Version: 1.0.0
+Summary: Azure Entra ID authentication extension for Python database drivers
+Author-email: Arjun Narendra <v-anarendra@microsoft.com>
+License: MIT
+Project-URL: Homepage, https://github.com/v-anarendra_microsoft/entra-id-integration-for-drivers
+Project-URL: Issues, https://github.com/v-anarendra_microsoft/entra-id-integration-for-drivers/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: azure-identity>=1.13.0
+Requires-Dist: azure-core>=1.24.0
+Provides-Extra: psycopg3
+Requires-Dist: psycopg[binary]>=3.1.0; extra == "psycopg3"
+Requires-Dist: aiohttp>=3.8.0; extra == "psycopg3"
+Provides-Extra: psycopg2
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "psycopg2"
+Provides-Extra: sqlalchemy
+Requires-Dist: sqlalchemy>=2.0.0; extra == "sqlalchemy"
+Requires-Dist: aiohttp>=3.8.0; extra == "sqlalchemy"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
+Requires-Dist: python-dotenv>=1.0.0; extra == "dev"
+Requires-Dist: mypy~=1.15; extra == "dev"
+Requires-Dist: ruff>=0.8.0; extra == "dev"
+Requires-Dist: types-psycopg2>=2.9.0; extra == "dev"
+Requires-Dist: psycopg-pool>=3.1.0; extra == "dev"
+Requires-Dist: testcontainers[postgres]>=3.7.0; extra == "dev"
+Requires-Dist: psycopg[binary]>=3.1.0; extra == "dev"
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "dev"
+Requires-Dist: sqlalchemy>=2.0.0; extra == "dev"
+Provides-Extra: all
+Requires-Dist: psycopg[binary]>=3.1.0; extra == "all"
+Requires-Dist: aiohttp>=3.8.0; extra == "all"
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "all"
+Requires-Dist: sqlalchemy>=2.0.0; extra == "all"
+Requires-Dist: pytest>=7.0.0; extra == "all"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "all"
+Requires-Dist: python-dotenv>=1.0.0; extra == "all"
+Requires-Dist: mypy~=1.15; extra == "all"
+Requires-Dist: ruff>=0.8.0; extra == "all"
+Requires-Dist: types-psycopg2>=2.9.0; extra == "all"
+Requires-Dist: psycopg-pool>=3.1.0; extra == "all"
+Requires-Dist: testcontainers[postgres]>=3.7.0; extra == "all"
+
+# azurepg-entra: Azure Database for PostgreSQL Entra ID Authentication
+
+This package provides seamless Azure Entra ID authentication for Python database drivers connecting to Azure Database for PostgreSQL. It supports both legacy and modern PostgreSQL drivers with automatic token management and connection pooling.
+
+## Features
+
+- **🔐 Azure Entra ID Authentication**: Automatic token acquisition and refresh for secure database connections
+- **🔄 Multi-Driver Support**: Works with psycopg2, psycopg3, and SQLAlchemy
+- **⚡ Connection Pooling**: Built-in support for both synchronous and asynchronous connection pools
+- **🏗️ Clean Architecture**: Simple package structure with `azure_postgresql_auth.psycopg2`, `azure_postgresql_auth.psycopg3`, and `azure_postgresql_auth.sqlalchemy`
+- **🔄 Automatic Token Management**: Handles token acquisition, validation, and refresh automatically
+- **🌐 Cross-platform**: Works on Windows, Linux, and macOS
+- **📦 Flexible Installation**: Optional dependencies for different driver combinations
+
+## Installation
+
+### Basic Installation
+
+Install the core package (includes Azure Identity dependencies only):
+```bash
+pip install azurepg-entra
+```
+
+### Driver-Specific Installation
+
+Choose the installation option based on which PostgreSQL drivers you need:
+
+```bash
+# For psycopg3 (modern psycopg, recommended for new projects)
+pip install "azurepg-entra[psycopg3]"
+
+# For psycopg2 (legacy support)
+pip install "azurepg-entra[psycopg2]"
+
+# For SQLAlchemy with psycopg3 backend
+pip install "azurepg-entra[sqlalchemy]"
+
+# All database drivers combined
+pip install "azurepg-entra[drivers]"
+
+# Everything including development tools
+pip install "azurepg-entra[all]"
+```
+
+### Development Installation
+
+Install from source for development:
+```bash
+git clone https://github.com/v-anarendra_microsoft/entra-id-integration-for-drivers.git
+cd entra-id-integration-for-drivers/python
+
+# Install with all dependencies for development
+pip install -e ".[all]"
+
+# Or install specific driver combinations
+pip install -e ".[psycopg3,dev]"
+```
+
+## Configuration
+
+### Environment Variables
+
+The samples use environment variables to configure database connections.
+
+Copy `.env.example` into a `.env` file in the same directory and update the variables.
+```env
+POSTGRES_SERVER=<your-server.postgres.database.azure.com>
+POSTGRES_DATABASE=<your_database_name>
+```
+
+## Quick Start
+
+### Running the Samples
+
+The repository includes comprehensive working examples in the `samples/` directory:
+
+- **`samples/psycopg2/getting_started/`**: psycopg2 (legacy driver support)
+- **`samples/psycopg3/getting_started/`**: psycopg3 examples (modern driver, recommended)
+- **`samples/sqlalchemy/getting_started/`**: SQLAlchemy examples with psycopg3 backend
+
+Configure your environment variables first, then run the samples:
+
+```bash
+# Copy and configure environment
+cp samples/psycopg3/getting_started/.env.example samples/psycopg3/getting_started/.env
+# Edit .env with your Azure PostgreSQL server details
+
+# Test psycopg2 (legacy driver)
+python samples/psycopg2/getting_started/create_db_connection_psycopg2.py --mode both
+
+# Test psycopg3 (modern driver, recommended)
+python samples/psycopg3/getting_started/create_db_connection_psycopg.py --mode both
+
+# Test SQLAlchemy 
+python samples/sqlalchemy/getting_started/create_db_connection_sqlalchemy.py --mode both
+```
+
+## Usage
+
+Choose the driver that best fits your project needs:
+
+- **psycopg3**: Modern PostgreSQL driver (recommended for new projects)
+- **psycopg2**: Legacy PostgreSQL driver (for existing projects)  
+- **SQLAlchemy**: High-level ORM/Core interface
+
+---
+
+## psycopg2 Driver (Legacy Support)
+
+> **Note**: psycopg2 is in maintenance mode. For new projects, consider using psycopg3 instead.
+
+The psycopg2 integration provides synchronous connection support with Azure Entra ID authentication through connection pooling.
+
+### Installation
+```bash
+pip install "azurepg-entra[psycopg2]"
+```
+
+### Connection Pooling (Recommended)
+
+```python
+from azure_postgresql_auth.psycopg2 import EntraConnection # import library
+from psycopg2 import pool # import to use pooling
+
+with pool.ThreadedConnectionPool(
+    minconn=1,
+    maxconn=5,
+    host="your-server.postgres.database.azure.com",
+    database="your_database",
+    connection_factory=EntraConnection
+) as connection_pool:
+```
+
+### Direct Connection
+
+```python
+from azure_postgresql_auth.psycopg2 import EntraConnection # import library
+
+with EntraConnection(
+    "postgresql://your-server.postgres.database.azure.com:5432/your_database"
+) as conn
+```
+
+---
+
+## psycopg3 Driver (Recommended)
+
+psycopg3 is the modern, actively developed PostgreSQL driver with native async support and better performance.
+
+### Installation
+```bash
+pip install "azurepg-entra[psycopg3]"
+```
+
+### Synchronous Connection
+
+```python
+from azure_postgresql_auth.psycopg3 import EntraConnection # import library
+from psycopg_pool import ConnectionPool # import to use pooling
+
+with ConnectionPool(
+    conninfo="postgresql://your-server.postgres.database.azure.com:5432/your_database",
+    connection_class=EntraConnection,
+    min_size=1,   # keep at least 1 connection always open
+    max_size=5,   # allow up to 5 concurrent connections
+) as pool
+```
+
+### Asynchronous Connection
+
+```python
+from azure_postgresql_auth.psycopg3 import AsyncEntraConnection # import library
+from psycopg_pool import AsyncConnectionPool # import to use pooling 
+
+async with AsyncConnectionPool(
+    conninfo="postgresql://your-server.postgres.database.azure.com:5432/your_database",
+    connection_class=AsyncEntraConnection,
+    min_size=1,   # keep at least 1 connection always open
+    max_size=5,   # allow up to 5 concurrent connections
+) as pool
+```
+
+---
+
+## SQLAlchemy Integration
+
+SQLAlchemy integration uses psycopg3 as the backend driver with automatic Entra ID authentication through event listeners.
+
+> **For more information**: See SQLAlchemy's documentation on [controlling how parameters are passed to the DBAPI connect function](https://docs.sqlalchemy.org/en/20/core/engines.html#controlling-how-parameters-are-passed-to-the-dbapi-connect-function).
+
+### Installation
+```bash
+pip install "azurepg-entra[sqlalchemy]"
+```
+
+### Synchronous Engine
+
+```python
+from sqlalchemy import create_engine
+from azure_postgresql_auth.sqlalchemy import enable_entra_authentication # import library
+
+with create_engine("postgresql+psycopg://your-server.postgres.database.azure.com/your_database") as engine:
+    # Enable Entra ID authentication
+    enable_entra_authentication(engine)
+    
+    # Core usage
+    with engine.connect() as conn:
+
+    # ORM usage
+    from sqlalchemy.orm import sessionmaker
+    Session = sessionmaker(bind=engine)
+```
+
+### Asynchronous Engine
+
+```python
+from sqlalchemy.ext.asyncio import create_async_engine
+from azure_postgresql_auth.sqlalchemy import enable_entra_authentication_async # import library
+
+async with create_async_engine("postgresql+psycopg://your-server.postgres.database.azure.com/your_database") as engine:
+    # Enable Entra ID authentication for async
+    enable_entra_authentication_async(engine)
+    
+    # Async Core usage
+    async with engine.connect() as conn:
+    
+    # Async ORM usage
+    from sqlalchemy.ext.asyncio import async_sessionmaker
+    AsyncSession = async_sessionmaker(engine, expire_on_commit=False)
+```
+
+## How It Works
+
+### Authentication Flow
+
+1. **Token Acquisition**: Uses Azure Identity libraries (`DefaultAzureCredential` by default) to acquire access tokens from Azure Entra ID
+2. **Automatic Refresh**: Tokens are automatically refreshed before each new database connection  
+3. **Secure Transport**: Tokens are passed as passwords in PostgreSQL connection strings over SSL
+4. **Server Validation**: Azure Database for PostgreSQL validates the token and establishes the authenticated connection
+5. **User Mapping**: The token's user principal name (UPN) is mapped to a PostgreSQL user for authorization
+
+### Token Scopes
+
+The package automatically requests the correct OAuth2 scopes:
+- **Database scope**: `https://ossrdbms-aad.database.windows.net/.default` (primary)
+- **Management scope**: `https://management.azure.com/.default` (fallback for managed identities)
+
+### Security Features
+
+- **🔒 Token-based authentication**: No passwords stored or transmitted
+- **⏰ Automatic expiration**: Tokens expire and are refreshed automatically
+- **🛡️ SSL enforcement**: All connections require SSL encryption
+- **🔑 Principle of least privilege**: Only database-specific scopes are requested
+---
+
+## Troubleshooting
+
+### Common Issues
+
+**Authentication Errors**
+```bash
+# Error: "password authentication failed"
+# Solution: Ensure your Azure identity has been granted access to the database
+# Run this SQL as a database administrator:
+CREATE ROLE "your-user@your-domain.com" WITH LOGIN;
+GRANT ALL PRIVILEGES ON DATABASE your_database TO "your-user@your-domain.com";
+```
+
+**Connection Timeouts**
+```python
+# Increase connection timeout for slow networks
+conn = SyncEntraConnection.connect(
+    "postgresql://server:5432/db", 
+    connect_timeout=30  # 30 seconds instead of default 10
+)
+```
+
+**Windows Async Issues**
+```python
+# Fix Windows event loop compatibility
+import asyncio
+import sys
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+```
+
+### Debug Logging
+
+Enable debug logging to troubleshoot authentication issues:
+
+```python
+import logging
+logging.basicConfig(level=logging.DEBUG)
+
+# This will show token acquisition and connection details
+conn = SyncEntraConnection.connect("postgresql://server:5432/db")
+```
+
+---
+
+## Contributing
+
+We welcome contributions! Please see [CONTRIBUTING.md](../CONTRIBUTING.md) for guidelines.
+
+
+---
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](../LICENSE) file for details.

azure_postgresql_auth-1.0.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+azure_postgresql_auth/__init__.py,sha256=2pRSes8cqoINx341xJCMrq6WmpUOs5QA3Nu3NPSaeBg,803
+azure_postgresql_auth/core.py,sha256=zW8XqcogluaXW9QB0ETS5Nfo-DdWPw6qrF4uXnz2nbs,9381
+azure_postgresql_auth/errors.py,sha256=rpu7SEJ6Uwldm4o48_7NTWsockJ_pBxe9cMLi-kD9BI,798
+azure_postgresql_auth/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+azure_postgresql_auth/psycopg2/__init__.py,sha256=twUVzJBoboQ45liZby_UOrStftS5C602cjdypp11dOY,620
+azure_postgresql_auth/psycopg2/entra_connection.py,sha256=Gm42vI4fWL6fdoPmUBDb95CK19amhHvjUFxoev-fFnE,3260
+azure_postgresql_auth/psycopg3/__init__.py,sha256=ChSOYtVhTzj2SF3svmhNwKVj0B7j_9CfoB5hZF_zsDg,815
+azure_postgresql_auth/psycopg3/async_entra_connection.py,sha256=jB4QIcMdkLIuxT9Y38yiNbn9oyb59yZZw-dRhLhASZQ,2972
+azure_postgresql_auth/psycopg3/entra_connection.py,sha256=cpJJ6AcfJMOv6eXIy8IGfCLvyrn99oG5IAVtd0R-R7g,2888
+azure_postgresql_auth/sqlalchemy/__init__.py,sha256=LttQhibGGG2rkaoBOOVjvpsYE4-szKT37PWoba42PS0,954
+azure_postgresql_auth/sqlalchemy/async_entra_connection.py,sha256=L3RjsqsAuj8NaYQRhMaiJbkBnblVAA8cYMfvAQQiOC8,2924
+azure_postgresql_auth/sqlalchemy/entra_connection.py,sha256=BcX5M6R0KaHF6Sflzl1V96BmwmVYjB-m80W7c1s1ydc,2743
+azure_postgresql_auth-1.0.0.dist-info/METADATA,sha256=RhDK4cCjui_unpI_QDOVOtYJ4S_3PkwWNzr0AOzMwrQ,12500
+azure_postgresql_auth-1.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+azure_postgresql_auth-1.0.0.dist-info/top_level.txt,sha256=tfXHngCqGdo4mArX4yjg3NJ1CQ_uUEmwnHaY9rkPnIc,22
+azure_postgresql_auth-1.0.0.dist-info/RECORD,,

azure_postgresql_auth-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

azure_postgresql_auth-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+azure_postgresql_auth