FEAT: Entra id authentication support for AQE

mssql_python/async_query/_auth.py

@@ -0,0 +1,140 @@
+"""Entra authentication for individual py-core async connections.
+
+Only credential instances are cached, never connections or raw tokens. Interactive
+credentials consequently share a signed-in account within this process; multi-user
+applications should supply their own per-user token providers.
+"""
+
+import asyncio
+import threading
+import warnings
+from typing import TYPE_CHECKING, Any, Optional
+
+from ..auth import (
+    ServicePrincipalAuth,
+    acquire_raw_token_from_credential,
+    extract_auth_type,
+)
+from ..constants import _AuthInternal, _KEY_AUTHENTICATION
+from ..exceptions import InterfaceError, OperationalError
+from ..logging import logger
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+
+    from ..connection import TokenProvider
+
+
+_credential_cache: dict[tuple[str, Optional[str]], "TokenCredential"] = {}
+_credential_cache_lock = threading.Lock()
+
+
+def _get_credential(auth_type: str, client_id: Optional[str]) -> "TokenCredential":
+    try:
+        from azure.identity import (
+            DefaultAzureCredential,
+            DeviceCodeCredential,
+            InteractiveBrowserCredential,
+            ManagedIdentityCredential,
+        )
+    except ImportError as error:
+        raise OperationalError(
+            "Azure authentication libraries are not installed. Install azure-identity.",
+            str(error),
+        ) from error
+
+    credential_classes = {
+        _AuthInternal.DEFAULT: DefaultAzureCredential,
+        _AuthInternal.DEVICE_CODE: DeviceCodeCredential,
+        _AuthInternal.INTERACTIVE: InteractiveBrowserCredential,
+        _AuthInternal.MSI: ManagedIdentityCredential,
+    }
+    credential_class = credential_classes[auth_type]
+    key = (auth_type, client_id)
+    with _credential_cache_lock:
+        credential = _credential_cache.get(key)
+        if credential is None:
+            try:
+                # Azure Identity's default interactive behavior is silent-first,
+                # prompting only when its cached account cannot supply a token.
+                credential = (
+                    credential_class(client_id=client_id) if client_id else credential_class()
+                )
+            except Exception as error:
+                logger.error(
+                    "Async authentication: credential initialization failed for %s",
+                    credential_class.__name__,
+                )
+                raise OperationalError(
+                    f"Failed to create {credential_class.__name__}",
+                    str(error),
+                ) from error
+            _credential_cache[key] = credential
+    return credential
+
+
+def _acquire_builtin_token(auth_type: str, client_id: Optional[str]) -> str:
+    credential = _get_credential(auth_type, client_id)
+    token, _ = acquire_raw_token_from_credential(credential)
+    return token
+
+
+async def prepare_async_authentication(
+    context: dict[str, Any],
+    token_provider: Optional["TokenProvider"] = None,
+) -> dict[str, Any]:
+    """Return a TDS login context without blocking the Python event loop.
+
+    ``token_provider`` accepts synchronous credentials only. Cancellation stops
+    connection setup but cannot interrupt an already-running credential call in
+    a worker thread. Service Principal callbacks are offloaded by py-core itself.
+    """
+    authenticated = dict(context)
+    auth_type = extract_auth_type({_KEY_AUTHENTICATION: context.get("authentication", "")})
+
+    if token_provider is not None:
+        if "authentication" in context:
+            raise InterfaceError(
+                "Cannot specify both 'token_provider' and 'Authentication'. Use one token source.",
+                "",
+            )
+        if not callable(getattr(token_provider, "get_token", None)):
+            raise InterfaceError(
+                "token_provider must have a callable .get_token(scope) method.", ""
+            )
+        dropped = [
+            name for name in ("user_name", "password", "trusted_connection") if name in context
+        ]
+        if dropped:
+            warnings.warn(
+                "token_provider is set, so the following connection credentials "
+                f"are ignored: {', '.join(dropped)}.",
+                UserWarning,
+                stacklevel=3,
+            )
+        raw_token, _ = await asyncio.to_thread(acquire_raw_token_from_credential, token_provider)
+    elif auth_type == _AuthInternal.SERVICE_PRINCIPAL:
+        try:
+            authenticated["entra_id_token_factory"] = ServicePrincipalAuth.make_token_factory(
+                context.get("user_name", ""), context.get("password", "")
+            )
+        except ValueError as error:
+            raise InterfaceError(
+                "ActiveDirectoryServicePrincipal requires UID (client_id) and PWD (client_secret).",
+                str(error),
+            ) from error
+        # py-core resolves the callback's authentication method from these
+        # fields, then clears the credentials from its native ClientContext.
+        return authenticated
+    elif auth_type:
+        client_id = None
+        if auth_type == _AuthInternal.MSI:
+            client_id = (context.get("user_name") or "").strip() or None
+        raw_token = await asyncio.to_thread(_acquire_builtin_token, auth_type, client_id)
+    else:
+        return authenticated
+
+    for key in ("authentication", "user_name", "password", "trusted_connection"):
+        authenticated.pop(key, None)
+    authenticated["access_token"] = raw_token
+    return authenticated

mssql_python/async_query/_connection_context.py

@@ -2,7 +2,9 @@
 
 from typing import Any
 
+from ..auth import extract_auth_type
 from ..connection_string_parser import _ConnectionStringParser
+from ..constants import _KEY_AUTHENTICATION
 from ..exceptions import InterfaceError, NotSupportedError
 from ..helpers import connstr_to_pycore_params
 
@@ -27,10 +29,15 @@ def build_async_connection_context(connection_str: str, timeout: int) -> dict[st
         raise ValueError("SERVER parameter is required in connection string")
 
     authentication = params.get("authentication", "").strip()
-    if authentication and authentication.lower() != "sqlpassword":
+    if (
+        authentication
+        and authentication.lower() != "sqlpassword"
+        and not extract_auth_type({_KEY_AUTHENTICATION: authentication})
+    ):
         raise NotSupportedError(
             f"Authentication mode '{authentication}' is not supported for async queries",
-            "Use Authentication=SqlPassword or Trusted_Connection for async queries",
+            "Use SqlPassword, Trusted_Connection, a supported ActiveDirectory mode, "
+            "or a token_provider.",
         )
 
     context = connstr_to_pycore_params(params, strict=True)

mssql_python/async_query/async_connection.py

@@ -6,9 +6,10 @@
     may change without notice.
 """
 
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from ..logging import logger
+from ._auth import prepare_async_authentication
 from ._native import load_py_core
 from ._connection_context import build_async_connection_context
 from .async_cursor import _AsyncCursor  # pyright: ignore[reportPrivateUsage]
@@ -26,6 +27,9 @@
     translate_py_core_exceptions,
 )
 
+if TYPE_CHECKING:
+    from ..connection import TokenProvider
+
 
 class _AsyncConnection:
     """Thin Python wrapper over ``mssql_py_core.PyAsyncConnection``.
@@ -56,8 +60,34 @@ async def connect(
         autocommit: bool = False,
         timeout: int = 0,
         python_logger: Optional[Any] = None,
+        *,
+        token_provider: Optional["TokenProvider"] = None,
     ) -> "_AsyncConnection":
-        """Establish an asynchronous connection from an ODBC connection string."""
+        """Establish a TDS connection using connection-string options.
+
+        Entra modes: ActiveDirectoryDefault, ActiveDirectoryMsi (optional UID
+        selects a user-assigned identity), ActiveDirectoryDeviceCode,
+        ActiveDirectoryInteractive, and client-secret ActiveDirectoryServicePrincipal.
+        Interactive authentication uses Python credentials on every platform.
+
+        Alternatively, supply a synchronous credential as ``token_provider``.
+        It must expose ``get_token(scope)`` returning an object with a non-empty
+        string ``token`` attribute; it cannot be combined with Authentication.
+        Async credentials (azure.identity.aio) are not accepted.
+
+        Token acquisition runs off the event loop. Each connect opens a new
+        physical connection; there is no connection pool, per-query refresh,
+        automatic query replay, or background refresh. Built-in credentials are
+        cached and may serve cached tokens. Interactive credentials share one
+        account per process; use per-user providers for multi-user applications.
+
+        The pre-acquired-token scope targets Azure commercial SQL. Service
+        Principal also currently uses the public-cloud authority. Sovereign
+        clouds and CAE claims challenges are not supported by these helpers.
+        ``timeout`` bounds native login, not preceding token acquisition.
+        Cancelling acquisition cannot stop a credential call already running in
+        a worker thread, but prevents subsequent native connection creation.
+        """
         logger_bridge = python_logger
         if logger_bridge is None and logger.is_debug_enabled:
             logger_bridge = logger
@@ -73,6 +103,9 @@ async def connect(
         with translate_py_core_exceptions():
             client_context_dict = build_async_connection_context(connection_str, timeout)
             py_core = load_py_core()
+            client_context_dict = await prepare_async_authentication(
+                client_context_dict, token_provider
+            )
             py_core_async_connection = await py_core.PyAsyncConnection.connect(
                 client_context_dict,
                 python_logger=logger_bridge,

tests/AsyncTest/test_002_async_connection.py

@@ -233,15 +233,16 @@ def test_bcp_conversion_does_not_fall_through_invalid_first_synonym():
     assert context == {}
 
 
-def test_async_connection_rejects_entra_authentication():
-    with pytest.raises(
-        NotSupportedError,
-        match="Authentication mode 'ActiveDirectoryDefault' is not supported",
-    ):
-        build_async_connection_context(
-            "Server=test-server.example.invalid;Authentication=ActiveDirectoryDefault",
-            0,
-        )
+def test_async_connection_context_accepts_entra_authentication_without_acquiring_token():
+    context = build_async_connection_context(
+        "Server=test-server.example.invalid;Authentication=ActiveDirectoryDefault",
+        0,
+    )
+
+    assert context == {
+        "server": "test-server.example.invalid",
+        "authentication": "ActiveDirectoryDefault",
+    }
 
 
 def test_async_connection_reports_invalid_authentication_mode():