From 72456860ac750ef4af716e12e2e8f46246f086a3 Mon Sep 17 00:00:00 2001 From: Pablo Ameijeiras Canay Date: Thu, 30 Jul 2026 16:48:43 +0200 Subject: [PATCH 1/2] feat(azure-postgresql-auth): add asyncpg SQLAlchemy engine factory --- .../azure-postgresql-auth/CHANGELOG.md | 6 ++ .../azure-postgresql-auth/README.md | 26 +++++++ .../sqlalchemy/__init__.py | 2 + .../sqlalchemy/asyncpg.py | 78 +++++++++++++++++++ .../dev_requirements.txt | 1 + .../azure-postgresql-auth/pyproject.toml | 5 +- .../tests/test_sqlalchemy_async.py | 70 +++++++++++++++-- 7 files changed, 181 insertions(+), 7 deletions(-) create mode 100644 sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py diff --git a/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md b/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md index bf401dcbdfe8..07e5bdc68430 100644 --- a/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md +++ b/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md @@ -1,5 +1,11 @@ # Release History +## Unreleased + +### Features Added + +- Added `create_asyncpg_engine` for SQLAlchemy async engines using `asyncpg` and asynchronous Microsoft Entra credentials. + ## 1.0.2 (2026-04-28) ### Bugs Fixed diff --git a/sdk/postgresql/azure-postgresql-auth/README.md b/sdk/postgresql/azure-postgresql-auth/README.md index d6517bdbeea1..d9d09cf7b0db 100644 --- a/sdk/postgresql/azure-postgresql-auth/README.md +++ b/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 diff --git a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/__init__.py b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/__init__.py index 029ad6487851..b71338012ed1 100644 --- a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/__init__.py +++ b/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", ] diff --git a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py new file mode 100644 index 000000000000..7aeddb90a0c5 --- /dev/null +++ b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py @@ -0,0 +1,78 @@ +# ------------------------------------------------------------------------- +# 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. + :param credential: Credential used to acquire Microsoft Entra access tokens. + :param kwargs: Keyword arguments forwarded to ``create_async_engine``. Values + supplied through ``connect_args`` are forwarded to ``asyncpg.connect``. + :return: An asynchronous SQLAlchemy engine. + :raises ~azure_postgresql_auth.CredentialValueError: If ``credential`` is not an + ``AsyncTokenCredential``. + :raises ImportError: If ``asyncpg`` is not installed. + """ + 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) diff --git a/sdk/postgresql/azure-postgresql-auth/dev_requirements.txt b/sdk/postgresql/azure-postgresql-auth/dev_requirements.txt index d26cf14fc463..77a51a088cd4 100644 --- a/sdk/postgresql/azure-postgresql-auth/dev_requirements.txt +++ b/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 diff --git a/sdk/postgresql/azure-postgresql-auth/pyproject.toml b/sdk/postgresql/azure-postgresql-auth/pyproject.toml index c964545381d7..dc6ce3a9c311 100644 --- a/sdk/postgresql/azure-postgresql-auth/pyproject.toml +++ b/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] diff --git a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py b/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py index 4e479e275c84..11233b2ce44d 100644 --- a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py +++ b/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: @@ -139,6 +144,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.""" From 0289fa923c38684d25054af7ff7202a5da753194 Mon Sep 17 00:00:00 2001 From: Pablo Ameijeiras Canay Date: Tue, 4 Aug 2026 10:02:57 +0200 Subject: [PATCH 2/2] Fix asyncpg factory pylint documentation --- .../azure_postgresql_auth/sqlalchemy/asyncpg.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py index 7aeddb90a0c5..0a8a3d7c7c8c 100644 --- a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py +++ b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/asyncpg.py @@ -35,13 +35,17 @@ def create_asyncpg_engine( 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. - :param kwargs: Keyword arguments forwarded to ``create_async_engine``. Values - supplied through ``connect_args`` are forwarded to ``asyncpg.connect``. + :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")