Add asyncpg SQLAlchemy engine factory for Entra authentication

sdk/postgresql/azure-postgresql-auth/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Features Added
+
+- Added `create_asyncpg_engine` for SQLAlchemy async engines using `asyncpg` and asynchronous Microsoft Entra credentials.
+
 ### Bugs Fixed
 
 - Fixed SQLAlchemy connection pools failing to create subsequent connections when using Entra authentication.

sdk/postgresql/azure-postgresql-auth/README.md

@@ -34,6 +34,9 @@ pip install "azure-postgresql-auth[psycopg2]"
 
 # For SQLAlchemy
 pip install "azure-postgresql-auth[sqlalchemy]"
+
+# For SQLAlchemy with asyncpg
+pip install "azure-postgresql-auth[sqlalchemy,asyncpg]"
 ```
 
 Install Azure Identity for credential support:
@@ -202,6 +205,29 @@ async with engine.connect() as conn:
     result = await conn.execute(text("SELECT 1"))
 ```
 
+### SQLAlchemy — Asynchronous asyncpg engine
+
+Use `create_asyncpg_engine` when the application uses `asyncpg` and an
+asynchronous Azure Identity credential. The helper acquires Entra tokens without
+blocking the event loop whenever the SQLAlchemy pool opens a physical connection.
+
+```python
+from azure.identity.aio import DefaultAzureCredential
+from azure_postgresql_auth.sqlalchemy import create_asyncpg_engine
+
+credential = DefaultAzureCredential()
+engine = create_asyncpg_engine(
+    "postgresql+asyncpg://your-server.postgres.database.azure.com/your_database?sslmode=require",
+    credential,
+)
+
+async with engine.connect() as conn:
+    result = await conn.execute(text("SELECT 1"))
+
+await engine.dispose()
+await credential.close()
+```
+
 ## Troubleshooting
 
 ### Authentication errors

sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/__init__.py

@@ -23,9 +23,11 @@
 """
 
 from .async_entra_connection import enable_entra_authentication_async
+from .asyncpg import create_asyncpg_engine
 from .entra_connection import enable_entra_authentication
 
 __all__ = [
     "enable_entra_authentication",
     "enable_entra_authentication_async",
+    "create_asyncpg_engine",
 ]

sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py

@@ -0,0 +1,82 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+
+"""asyncpg integration for SQLAlchemy asynchronous engines."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from azure.core.credentials_async import AsyncTokenCredential
+
+try:
+    from sqlalchemy.engine import URL, make_url
+    from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+except ImportError as e:
+    raise ImportError(
+        "SQLAlchemy dependencies are not installed. Install them with: pip install azure-postgresql-auth[sqlalchemy]"
+    ) from e
+
+from azure_postgresql_auth.core import get_entra_conninfo_async
+from azure_postgresql_auth.errors import CredentialValueError, EntraConnectionValueError
+
+
+def create_asyncpg_engine(
+    url: str | URL,
+    credential: AsyncTokenCredential,
+    **kwargs: Any,
+) -> AsyncEngine:
+    """Create an asyncpg SQLAlchemy engine authenticated with Microsoft Entra ID.
+
+    The returned engine obtains Entra connection information asynchronously whenever
+    SQLAlchemy creates a physical connection for its pool.
+
+    :param url: SQLAlchemy URL using the ``postgresql+asyncpg`` dialect.
+    :type url: str or ~sqlalchemy.engine.URL
+    :param credential: Credential used to acquire Microsoft Entra access tokens.
+    :type credential: ~azure.core.credentials_async.AsyncTokenCredential
+    :return: An asynchronous SQLAlchemy engine.
+    :rtype: ~sqlalchemy.ext.asyncio.AsyncEngine
+    :raises ~azure_postgresql_auth.CredentialValueError: If ``credential`` is not an
+        ``AsyncTokenCredential``.
+    :raises ImportError: If ``asyncpg`` is not installed.
+
+    Additional keyword arguments are forwarded to ``create_async_engine``. Values
+    supplied through ``connect_args`` are forwarded to ``asyncpg.connect``.
+    """
+    if not isinstance(credential, AsyncTokenCredential):
+        raise CredentialValueError("credential is required and must be an AsyncTokenCredential for asyncpg")
+
+    try:
+        import asyncpg
+    except ImportError as e:
+        raise ImportError(
+            "asyncpg dependencies are not installed. Install them with: pip install azure-postgresql-auth[asyncpg]"
+        ) from e
+
+    parsed_url = make_url(url)
+    connect_args = parsed_url.translate_connect_args(username="user", database="database")
+    connect_args.update(dict(parsed_url.query))
+    connect_args.update(kwargs.pop("connect_args", {}))
+
+    if "sslmode" in connect_args and "ssl" not in connect_args:
+        connect_args["ssl"] = connect_args["sslmode"]
+    connect_args.pop("sslmode", None)
+
+    async def async_creator() -> Any:
+        try:
+            entra_conninfo = await get_entra_conninfo_async(credential)
+        except Exception as e:
+            raise EntraConnectionValueError("Could not retrieve Entra credentials") from e
+
+        connection_kwargs = {
+            **connect_args,
+            "user": connect_args.get("user", entra_conninfo["user"]),
+            "password": connect_args.get("password", entra_conninfo["password"]),
+        }
+        return await asyncpg.connect(**connection_kwargs)
+
+    return create_async_engine(parsed_url, async_creator=async_creator, **kwargs)

sdk/postgresql/azure-postgresql-auth/dev_requirements.txt

@@ -2,6 +2,7 @@
 ../../core/azure-core
 ../../identity/azure-identity
 aiohttp
+asyncpg>=0.29.0
 pytest
 pytest-asyncio
 psycopg2-binary>=2.9.0

sdk/postgresql/azure-postgresql-auth/pyproject.toml

@@ -37,7 +37,10 @@ psycopg2 = [
     "psycopg2-binary>=2.9.0",
 ]
 sqlalchemy = [
-    "sqlalchemy>=2.0.0",
+    "sqlalchemy[asyncio]>=2.0.0",
+]
+asyncpg = [
+    "asyncpg>=0.29.0",
 ]
 
 [project.urls]

sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py

@@ -9,18 +9,23 @@
 from __future__ import annotations
 
 import asyncio
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from azure_postgresql_auth.errors import CredentialValueError, EntraConnectionValueError
+from azure_postgresql_auth.sqlalchemy import create_asyncpg_engine, enable_entra_authentication_async
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import create_async_engine
-
-from azure_postgresql_auth.errors import CredentialValueError, EntraConnectionValueError
-from azure_postgresql_auth.sqlalchemy import enable_entra_authentication_async
-
-from utils import TEST_USERS, MockTokenCredential, capture_event_handler, create_valid_jwt_token
+from utils import (
+    TEST_USERS,
+    MockAsyncTokenCredential,
+    MockTokenCredential,
+    capture_event_handler,
+    create_valid_jwt_token,
+)
 
 ASYNC_MODULE = "azure_postgresql_auth.sqlalchemy.async_entra_connection"
+ASYNCPG_MODULE = "azure_postgresql_auth.sqlalchemy.asyncpg"
 
 
 class TestSqlalchemyAsyncEntraAuthentication:
@@ -145,6 +150,59 @@ def test_entra_credential_failure_raises_error_async(self, mock_get_conninfo):
             handler(MagicMock(), MagicMock(), [], {"credential": credential})
 
 
+class TestSqlalchemyAsyncpgEntraAuthentication:
+    """Tests for asyncpg SQLAlchemy Entra authentication integration."""
+
+    @patch(f"{ASYNCPG_MODULE}.create_async_engine")
+    @patch(f"{ASYNCPG_MODULE}.get_entra_conninfo_async", new_callable=AsyncMock)
+    @patch("asyncpg.connect", new_callable=AsyncMock)
+    @pytest.mark.asyncio
+    async def test_creates_asyncpg_connection_with_entra_credentials(
+        self, mock_connect, mock_get_conninfo, mock_create_engine
+    ):
+        """Test that the async creator injects Entra credentials into asyncpg."""
+        token = create_valid_jwt_token(TEST_USERS["ENTRA_USER"])
+        credential = MockAsyncTokenCredential(token)
+        mock_get_conninfo.return_value = {"user": TEST_USERS["ENTRA_USER"], "password": token}
+
+        create_asyncpg_engine(
+            "postgresql+asyncpg://server.example:5432/database?sslmode=require",
+            credential,
+        )
+
+        async_creator = mock_create_engine.call_args.kwargs["async_creator"]
+        await async_creator()
+
+        mock_get_conninfo.assert_awaited_once_with(credential)
+        mock_connect.assert_awaited_once_with(
+            host="server.example",
+            port=5432,
+            database="database",
+            user=TEST_USERS["ENTRA_USER"],
+            password=token,
+            ssl="require",
+        )
+
+    def test_invalid_asyncpg_credential_raises_error(self):
+        """Test that a synchronous credential is rejected before creating an engine."""
+        with pytest.raises(CredentialValueError, match="AsyncTokenCredential"):
+            create_asyncpg_engine("postgresql+asyncpg://server.example/database", MagicMock())
+
+    @patch(f"{ASYNCPG_MODULE}.create_async_engine")
+    @patch(f"{ASYNCPG_MODULE}.get_entra_conninfo_async", new_callable=AsyncMock)
+    @pytest.mark.asyncio
+    async def test_asyncpg_credential_failure_raises_connection_error(self, mock_get_conninfo, mock_create_engine):
+        """Test that async credential failures are surfaced as Entra connection errors."""
+        mock_get_conninfo.side_effect = Exception("auth failed")
+        credential = MockAsyncTokenCredential(create_valid_jwt_token(TEST_USERS["ENTRA_USER"]))
+
+        create_asyncpg_engine("postgresql+asyncpg://server.example/database", credential)
+
+        async_creator = mock_create_engine.call_args.kwargs["async_creator"]
+        with pytest.raises(EntraConnectionValueError, match="Could not retrieve Entra credentials"):
+            await async_creator()
+
+
 @pytest.mark.live_test_only
 class TestSqlalchemyAsyncEntraAuthenticationLive:
     """Live tests for asynchronous SQLAlchemy with enable_entra_authentication_async."""