Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions sdk/postgresql/azure-postgresql-auth/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down
26 changes: 26 additions & 0 deletions sdk/postgresql/azure-postgresql-auth/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Comment thread
mattboentoro marked this conversation as resolved.

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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@
"""

from .async_entra_connection import enable_entra_authentication_async
from .asyncpg import create_asyncpg_engine
Comment thread
mattboentoro marked this conversation as resolved.
from .entra_connection import enable_entra_authentication

__all__ = [
"enable_entra_authentication",
"enable_entra_authentication_async",
"create_asyncpg_engine",
]
Original file line number Diff line number Diff line change
@@ -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"]),
}
Comment thread
mattboentoro marked this conversation as resolved.
Comment thread
mattboentoro marked this conversation as resolved.
return await asyncpg.connect(**connection_kwargs)

return create_async_engine(parsed_url, async_creator=async_creator, **kwargs)
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
../../core/azure-core
../../identity/azure-identity
aiohttp
asyncpg>=0.29.0
pytest
pytest-asyncio
psycopg2-binary>=2.9.0
Expand Down
5 changes: 4 additions & 1 deletion sdk/postgresql/azure-postgresql-auth/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,10 @@ psycopg2 = [
"psycopg2-binary>=2.9.0",
]
sqlalchemy = [
"sqlalchemy>=2.0.0",
"sqlalchemy[asyncio]>=2.0.0",
Comment thread
mattboentoro marked this conversation as resolved.
]
asyncpg = [
"asyncpg>=0.29.0",
]

[project.urls]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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."""
Comment thread
mattboentoro marked this conversation as resolved.

@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."""
Expand Down
Loading