arrowsql 0.1.0__cp313-cp313-win_amd64.whl

arrowsql/_rust_auth.py

@@ -0,0 +1,170 @@
+"""
+Authentication= handling for the Rust core (RUST-PORT, new module).
+
+The Rust core logs in with exactly what the connection string says; it has no
+msodbcsql underneath that interprets ``Authentication=``. Upstream's auth.py
+returns ``None`` for the modes ODBC "handles natively" (ServicePrincipal,
+Windows Interactive) and does not know ``SqlPassword`` / ``ActiveDirectoryPassword``
+/ ``ActiveDirectoryIntegrated`` at all, so without this module those would fall
+through to a SQL login with the client id or UPN. This module keeps auth.py
+(and the tests that pin its return values) unchanged and makes every
+``Authentication`` value either work or fail with a clear error.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Callable, Dict, Optional
+
+from arrowsql.constants import (
+    AuthType,
+    _KEY_AUTHENTICATION,
+    _KEY_PWD,
+    _KEY_UID,
+)
+from arrowsql.exceptions import InterfaceError
+
+_SQL_PASSWORD = "sqlpassword"
+_AD_PASSWORD = "activedirectorypassword"
+_MANAGED_IDENTITY = "activedirectorymanagedidentity"
+
+# Acquired by Connection.__init__ (auth.py) before the native connect.
+_TOKEN_MODES = frozenset(
+    {
+        AuthType.INTERACTIVE.value,
+        AuthType.DEVICE_CODE.value,
+        AuthType.DEFAULT.value,
+        AuthType.MSI.value,
+    }
+)
+_UNSUPPORTED = {
+    "activedirectoryintegrated": (
+        "Authentication=ActiveDirectoryIntegrated is not supported by the Rust core. "
+        "For Kerberos (GSSAPI) integrated login use Trusted_Connection=yes; for Entra ID "
+        "pass a token with token_provider= or Authentication=ActiveDirectoryDefault."
+    ),
+    "activedirectoryworkloadidentity": (
+        "Authentication=ActiveDirectoryWorkloadIdentity is not supported by the Rust core. "
+        "Pass a token with token_provider= or use Authentication=ActiveDirectoryDefault."
+    ),
+}
+_KNOWN = (
+    _TOKEN_MODES
+    | {AuthType.SERVICE_PRINCIPAL.value, _SQL_PASSWORD, _AD_PASSWORD, _MANAGED_IDENTITY}
+    | set(_UNSUPPORTED)
+)
+
+# Public client id of the Azure CLI, the one tenant-agnostic first-party client
+# that allows the username/password (ROPC) grant without an app registration.
+_PASSWORD_CLIENT_ID = "04b07795-8ddb-461a-bbee-02f9e1bf7b46"
+
+_AZURE_HINT = "Install the optional dependency: pip install arrowsql[azure]"
+
+
+def _require_azure_identity(mode: str) -> None:
+    try:
+        import azure.identity  # noqa: F401  pylint: disable=import-outside-toplevel,unused-import
+    except ImportError as e:
+        raise InterfaceError(
+            driver_error=f"Authentication={mode} needs azure-identity. {_AZURE_HINT}",
+            ddbc_error="azure-identity is not installed.",
+        ) from e
+
+
+def make_password_token_factory(username: str, password: str) -> Callable[[str, str, str], bytes]:
+    """Token factory for ``ActiveDirectoryPassword``.
+
+    Same contract as ``ServicePrincipalAuth.make_token_factory``: called by the
+    native login with ``(spn, sts_url, auth_method)`` once the server has named
+    the tenant, returns the JWT as bare UTF-16LE bytes.
+    """
+    if not username or not password:
+        raise ValueError("ActiveDirectoryPassword auth requires UID (user principal name) and PWD")
+
+    def _factory(spn: str, sts_url: str, _auth_method: str) -> bytes:
+        # pylint: disable=import-outside-toplevel
+        from azure.identity import UsernamePasswordCredential
+        from azure.core.exceptions import ClientAuthenticationError
+        from arrowsql.auth import _parse_tenant_id
+
+        if not spn:
+            raise RuntimeError("ActiveDirectoryPassword: empty SPN from server")
+        tenant_id = _parse_tenant_id(sts_url)
+        if not tenant_id:
+            raise RuntimeError(f"Could not extract tenant_id from STS URL: {sts_url!r}")
+        scope = spn if spn.endswith("/.default") else spn.rstrip("/") + "/.default"
+        try:
+            credential = UsernamePasswordCredential(
+                client_id=_PASSWORD_CLIENT_ID,
+                username=username,
+                password=password,
+                tenant_id=tenant_id,
+            )
+            token = credential.get_token(scope).token
+        except ClientAuthenticationError:
+            # The provider text can echo the account; keep it out of the exception chain.
+            raise RuntimeError(
+                "ActiveDirectoryPassword authentication failed (wrong credentials, MFA "
+                "required, or the account is not allowed the password grant)"
+            ) from None
+        return token.encode("utf-16-le")
+
+    return _factory
+
+
+def prepare(parsed_params: Dict[str, str]) -> Optional[Callable[[str, str, str], bytes]]:
+    """Validate ``Authentication`` and return the federated token factory, if any.
+
+    ``ActiveDirectoryManagedIdentity`` is rewritten in *parsed_params* to its
+    ``ActiveDirectoryMSI`` alias so the token path in ``Connection.__init__``
+    picks it up. Raises ``InterfaceError`` for unknown or unsupported values.
+    """
+    raw = parsed_params.get(_KEY_AUTHENTICATION, "")
+    mode = raw.strip().lower()
+    if not mode:
+        return None
+    if mode not in _KNOWN:
+        raise InterfaceError(
+            driver_error=f"Invalid value specified for connection string attribute 'Authentication': {raw}",
+            ddbc_error=f"Unknown Authentication value: {raw}",
+        )
+    if mode in _UNSUPPORTED:
+        raise InterfaceError(driver_error=_UNSUPPORTED[mode], ddbc_error=f"Unsupported: {raw}")
+    if mode == _MANAGED_IDENTITY:
+        parsed_params[_KEY_AUTHENTICATION] = "ActiveDirectoryMSI"
+        return None
+    if mode not in (AuthType.SERVICE_PRINCIPAL.value, _AD_PASSWORD):
+        return None
+    _require_azure_identity(raw)
+    uid = (parsed_params.get(_KEY_UID) or "").strip()
+    pwd = parsed_params.get(_KEY_PWD) or ""
+    try:
+        if mode == _AD_PASSWORD:
+            return make_password_token_factory(uid, pwd)
+        # pylint: disable=import-outside-toplevel
+        from arrowsql.auth import ServicePrincipalAuth
+
+        return ServicePrincipalAuth.make_token_factory(uid, pwd)
+    except ValueError as e:
+        raise InterfaceError(
+            driver_error=f"Authentication={raw} requires UID and PWD: {e}",
+            ddbc_error=str(e),
+        ) from e
+
+
+def reject_unhandled(parsed_params: Dict[str, str]) -> None:
+    """Raise when a token mode produced no token (Windows Interactive).
+
+    Upstream leaves Windows Interactive to ODBC's native AAD prompt, which the
+    Rust core does not have; connecting anyway would be a SQL login.
+    """
+    mode = parsed_params.get(_KEY_AUTHENTICATION, "").strip().lower()
+    if mode in _TOKEN_MODES:
+        raise InterfaceError(
+            driver_error=(
+                f"Authentication={parsed_params[_KEY_AUTHENTICATION]} is not available on "
+                f"{sys.platform}: the Rust core has no native interactive flow. "
+                "Use token_provider= or Authentication=ActiveDirectoryDefault."
+            ),
+            ddbc_error="No native interactive Entra ID flow.",
+        )

arrowsql/_server_sqlstate.py

@@ -0,0 +1,11 @@
+"""
+SQL Server engine error number -> ODBC 3.x SQLSTATE.
+
+The table lives in the native layer (src/diag.rs); this module re-exports it
+for the async py-core path, so exceptions still carry the SQLSTATE that
+msodbcsql would report (#581).
+"""
+
+from arrowsql.ddbc_bindings import sqlstate_for_server_error
+
+__all__ = ["sqlstate_for_server_error"]

arrowsql/_utf8_strings.py

@@ -0,0 +1,78 @@
+"""RUST-PORT: UTF-8 transfer of national character columns for Arrow reads.
+
+SQL Server sends nvarchar/nchar/ntext as UTF-16 (two bytes per character,
+even ASCII) and TDS has no compression. ``utf8_query`` rewrites a query so
+those columns travel as ``varchar`` in a UTF-8 collation, about half the
+bytes for mostly-ASCII text; the decoded Arrow type is the same string type.
+The SQL text is built in Rust (``ddbc_bindings.utf8_wrap_query``); this module
+only runs the describe round trip and decides when to skip.
+
+Every failure path returns the original query: the rewrite is an
+optimization, never a reason for a read to fail.
+"""
+
+import re
+import weakref
+from contextlib import closing
+from typing import Any, Optional
+
+from arrowsql import ddbc_bindings
+from arrowsql.logging import logger
+
+# One round trip: server capability plus the described, visible columns.
+_DESCRIBE_SQL = (
+    "SELECT CAST(SERVERPROPERTY('ProductMajorVersion') AS int), "
+    "CAST(SERVERPROPERTY('EngineEdition') AS int), "
+    "d.name, d.system_type_id, d.max_length, d.is_nullable "
+    "FROM sys.dm_exec_describe_first_result_set(?, NULL, 0) AS d "
+    "WHERE d.is_hidden = 0 ORDER BY d.column_ordinal"
+)
+# Rewrite only when UTF-8 is at most this share of the UTF-16 bytes of the
+# sampled rows. CJK is 3 bytes per character in UTF-8 against 2 in UTF-16.
+_MAX_SIZE_RATIO = 0.9
+# The rewrite puts the query in a derived table, which drops ORDER BY.
+_ORDER_BY = re.compile(r"\bORDER\s+BY\b", re.IGNORECASE)
+
+# connection -> server accepts UTF-8 collations (learned from the first describe).
+_server_support: "weakref.WeakKeyDictionary[Any, bool]" = weakref.WeakKeyDictionary()
+
+
+def utf8_query(conn: Any, query: str) -> str:
+    """``query`` rewritten to transfer nvarchar/nchar/ntext columns as UTF-8.
+
+    Returns ``query`` itself when the server predates UTF-8 collations
+    (SQL Server 2019 / Azure SQL), the query has an ORDER BY, no column would
+    change, the column names are empty or duplicated, the describe or the
+    sample probe fails, or the first 256 rows are not at least 10% smaller in
+    UTF-8 (CJK text grows: 3 bytes per character against 2).
+    """
+    if _ORDER_BY.search(query) or _server_support.get(conn) is False:
+        return query
+    try:
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(_DESCRIBE_SQL, (query,))
+            rows = cursor.fetchall()
+            if not rows:
+                return query
+            supported = ddbc_bindings.utf8_collations_supported(rows[0][0], rows[0][1])
+            try:
+                _server_support[conn] = supported
+            except TypeError:
+                pass
+            if not supported:
+                return query
+            plan = ddbc_bindings.utf8_wrap_query(query, [(r[2], r[3], r[4], r[5]) for r in rows])
+            if plan is None:
+                return query
+            wrapped, sample_sql = plan
+            # Runs the query's first rows: also finds out that the rewrite is
+            # invalid (batches, EXEC, ...) and how big the strings are in UTF-8.
+            cursor.execute(sample_sql)
+            utf8_bytes, utf16_bytes = cursor.fetchone()
+            if not utf16_bytes or utf8_bytes > _MAX_SIZE_RATIO * utf16_bytes:
+                logger.debug("utf8_strings: no saving in sample (%s vs %s bytes)", utf8_bytes, utf16_bytes)
+                return query
+            return wrapped
+    except Exception as exc:
+        logger.debug("utf8_strings: keeping the original query (%s)", exc)
+        return query