python-rust-integration

sdk/cosmos/azure-cosmos/azure/cosmos/_backend/base.py

@@ -0,0 +1,120 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+"""Abstract backend type plus the two small data classes used to talk to it.
+
+Every concrete backend (``CorePythonBackend``, ``RustBackend``)
+implements the ``CosmosBackend`` ABC defined here. Container methods
+call the backend through this ABC so that the dispatch site never has
+to know which underlying transport actually ran the request.
+
+``PreparedRequest`` is what the helper hands a backend; ``BackendResponse``
+is what the helper expects back. Both are frozen dataclasses (immutable
+after construction) so a backend cannot accidentally mutate the input it
+was handed or the output it produced.
+"""
+from __future__ import annotations
+
+import abc
+from dataclasses import dataclass, field
+from typing import Any, Mapping, Optional
+
+from azure.core.utils import CaseInsensitiveDict
+
+
+@dataclass(frozen=True)
+class PreparedRequest:
+    """A single Cosmos operation, fully prepared and ready to send.
+
+    The helper layer that owns request preparation builds one of these
+    from the user's ``create_item`` arguments. Both backends receive
+    the *same* instance so neither one re-derives the wire format from
+    the original kwargs; that is the only way to guarantee byte-for-byte
+    parity between them.
+    """
+
+    #: e.g. ``"dbs/{db}/colls/{coll}"``.
+    container_link: str
+
+    #: The request body, already serialized to JSON bytes by the helper.
+    body_bytes: bytes
+
+    #: The partition-key header value, already serialized to its on-wire
+    #: JSON shape (e.g. ``'["customerA"]'`` for a single-value PK).
+    partition_key_header: str
+
+    #: Everything else that needs to ride on the request: triggers,
+    #: indexing directive, intended-collection-rid, etc.
+    headers: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class BackendResponse:
+    """Normalized shape every backend produces, regardless of transport.
+
+    The core-python backend builds one of these from an azure-core
+    ``HttpResponse``; the Rust backend builds one from the PyO3 return
+    value. Code above the backend never branches on which backend
+    handled the call — it just reads these fields.
+    """
+
+    #: HTTP status code returned by the service.
+    status_code: int
+
+    #: Cosmos sub-status code (header ``x-ms-substatus``); ``0`` if absent.
+    sub_status: int = 0
+
+    #: Full response header map. Populated from the underlying response
+    #: object so that long-tail headers (e.g. ``x-ms-cosmos-llsn``) survive
+    #: the trip across the backend boundary.
+    headers: Optional[CaseInsensitiveDict] = None
+
+    #: Raw response body bytes. May be empty for a 204 / no-content reply.
+    body: bytes = b""
+
+    #: Per-backend diagnostics blob the helper does not introspect. The
+    #: core-python backend stores its existing diagnostics here; the Rust
+    #: backend stores its structured diagnostics here.
+    diagnostics: Any = None
+
+
+class CosmosBackend(abc.ABC):
+    """Abstract dispatch target for any Cosmos operation (sync variant).
+
+    Every sync backend (``CorePythonBackend``, ``RustBackend``) inherits
+    from this class. The container's dispatch site holds one of these by
+    interface and calls ``create_item`` on it without knowing which
+    concrete backend it has.
+
+    Today, request preparation and response parsing still live in
+    their existing locations (``Container.create_item`` and the
+    ``CosmosClientConnection`` helpers), so a backend that returns
+    ``None`` from ``create_item`` is interpreted as "I have nothing to
+    return; the caller should run its existing in-place
+    implementation." Once the helper layer takes over request prep and
+    response parsing the contract tightens so that every backend
+    produces a real ``BackendResponse`` and the ``Optional`` annotations
+    go away.
+    """
+
+    #: Short identifier used in two places: the startup INFO log line
+    #: emitted at client construction, and the per-request user-agent
+    #: suffix (``backend=<name>``). Concrete subclasses set this from
+    #: ``constants.BACKEND_NAME_CORE_PYTHON`` or
+    #: ``constants.BACKEND_NAME_RUST``.
+    name: str = "abstract"
+
+    @abc.abstractmethod
+    def create_item(self, prepared: Optional[PreparedRequest]) -> Optional[BackendResponse]:
+        """Issue a single ``create_item`` call.
+
+        Returning ``None`` (the temporary contract until the helper
+        layer lands) tells the caller to use its existing in-place
+        implementation. Returning a ``BackendResponse`` tells the
+        caller to use that response directly. ``prepared`` may be
+        ``None`` while the caller still owns request preparation.
+        """
+        ...
+

sdk/cosmos/azure-cosmos/azure/cosmos/_backend/constants.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.
+# -------------------------------------------------------------------------
+"""Constants for the backend dispatch layer.
+
+Why this file exists:
+
+    The same strings — the two backend names, the env var, the per-request
+    kwarg key — show up in the factory, the two backend classes, the
+    user-agent policy, the container dispatch site, and the tests. If any
+    one of those copies drifts (for example, somebody types "corepython"
+    instead of "core-python"), the bug only shows up at runtime and is
+    annoying to track down.
+
+    Putting the strings in one place fixes that. Every other module that
+    needs one of these names imports it from here. Tests assert against
+    the same constants the production code uses, so renaming a value in
+    this file is a one-line change that the type checker and test suite
+    immediately tell you about.
+"""
+from __future__ import annotations
+
+# ---------------------------------------------------------------------------
+# Backend names
+# ---------------------------------------------------------------------------
+#
+# The two values that the user can set on the constructor kwarg ``_backend``
+# or in the ``COSMOS_BACKEND`` environment variable. They also show up in
+# the per-request user-agent suffix (``backend=core-python`` or
+# ``backend=rust``) so the server-side log knows which path handled a call.
+
+#: The default backend. Routes through the existing in-tree azure-core
+#: pipeline plus the ``CosmosClientConnection`` helpers. This is *not*
+#: deprecated; it is the production path today and stays that way until
+#: a future release decides the Rust path is ready to be the default.
+BACKEND_NAME_CORE_PYTHON = "core-python"
+
+#: The Rust-driver-backed path. Opt-in only today; the implementation
+#: itself lands incrementally in later changes. Selecting this today
+#: causes ``RustBackend.create_item`` to raise ``NotImplementedError``
+#: so the developer gets a loud signal that they are running against a
+#: stub.
+BACKEND_NAME_RUST = "rust"
+
+#: Tuple of every accepted backend name. The factory validates against
+#: this; an unknown value raises ``ValueError`` at client construction.
+VALID_BACKEND_NAMES = (BACKEND_NAME_CORE_PYTHON, BACKEND_NAME_RUST)
+
+#: The backend the factory picks when the caller did not pass ``_backend=``
+#: and the env var is unset. Stays ``core-python`` until a future release
+#: explicitly flips it.
+DEFAULT_BACKEND_NAME = BACKEND_NAME_CORE_PYTHON
+
+
+# ---------------------------------------------------------------------------
+# Selection knob
+# ---------------------------------------------------------------------------
+
+#: Environment variable read by the factory when no explicit ``_backend=``
+#: was passed to the client constructor. Lets a developer or operator
+#: switch backends without a code change. Precedence:
+#: constructor kwarg  >  this env var  >  ``DEFAULT_BACKEND_NAME``.
+BACKEND_ENV_VAR = "COSMOS_BACKEND"
+
+
+# ---------------------------------------------------------------------------
+# Per-request stamp
+# ---------------------------------------------------------------------------
+
+#: Key the container dispatch site writes into ``**kwargs`` (which become
+#: the azure-core request-context options) recording which backend
+#: actually handled a single call. ``CosmosUserAgentPolicy`` reads it from
+#: there and appends ``backend=<name>`` to the User-Agent header.
+#:
+#: The value carried under this key is one of ``BACKEND_NAME_CORE_PYTHON``
+#: or ``BACKEND_NAME_RUST``. It is stamped *after* the forced-fallback
+#: decision so a request that fell back to core-python on a Rust-default
+#: client correctly reports ``backend=core-python`` server-side.
+REQUEST_OPTION_BACKEND_KEY = "cosmos_backend"
+

sdk/cosmos/azure-cosmos/azure/cosmos/_backend/core_python.py

@@ -0,0 +1,42 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+"""Sync core-python backend.
+
+"Core-python" is the existing in-tree path: an azure-core pipeline plus
+the ``CosmosClientConnection`` helpers that have shipped with this
+package for years. It is the default and is *not* deprecated. It is also
+the always-available fallback the dispatch site uses for any request
+whose kwargs the Rust backend does not support yet (today: the two
+kwargs ``availability_strategy`` and ``retry_write``).
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+from .base import BackendResponse, CosmosBackend, PreparedRequest
+from .constants import BACKEND_NAME_CORE_PYTHON
+
+
+class CorePythonBackend(CosmosBackend):
+    """Routes ``create_item`` calls through the existing azure-core stack.
+
+    Today this class's ``create_item`` returns ``None`` to tell the
+    dispatch site "delegate to the existing in-place call." When the
+    request-preparation helper and the response-parsing helper land,
+    this method will build an azure-core ``HttpRequest`` from the
+    ``PreparedRequest``, send it through
+    ``client_connection.pipeline_client``, and return a populated
+    ``BackendResponse``.
+    """
+
+    name = BACKEND_NAME_CORE_PYTHON
+
+    def create_item(self, prepared: Optional[PreparedRequest]) -> Optional[BackendResponse]:
+        # Returning None tells the container's dispatch site to fall
+        # through to the existing in-place create_item code path. The
+        # real adapter lands once the helper layer exists.
+        return None
+

sdk/cosmos/azure-cosmos/azure/cosmos/_backend/factory.py

@@ -0,0 +1,74 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+"""Factory that picks which backend a single client will use.
+
+A ``CosmosClient`` calls ``make_backend(...)`` exactly once at
+construction time and stores the returned object. Every subsequent
+``create_item`` call on every container created by that client dispatches
+through the same instance.
+
+Selection precedence (highest wins):
+
+  1. The value the caller passed to the client constructor as the
+     private kwarg ``_backend=``.
+  2. The value of the ``COSMOS_BACKEND`` environment variable.
+  3. The default, ``core-python``.
+
+If a value comes in that is not ``"core-python"`` or ``"rust"``, the
+factory raises ``ValueError`` immediately at client construction time.
+Failing loud is intentional — a typo'd value silently falling back to
+the default would mask configuration drift in production.
+"""
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from .base import CosmosBackend
+from .constants import (
+    BACKEND_ENV_VAR,
+    BACKEND_NAME_RUST,
+    DEFAULT_BACKEND_NAME,
+    VALID_BACKEND_NAMES,
+)
+from .core_python import CorePythonBackend
+from .rust import RustBackend
+
+
+def resolve_backend_name(explicit: Optional[str]) -> str:
+    """Return one of the values in ``VALID_BACKEND_NAMES`` after applying
+    the precedence rules above.
+
+    Shared between the sync and async factories so the precedence rules,
+    valid values, and error message live in one place.
+    """
+    if explicit is not None:
+        choice = explicit
+    else:
+        choice = os.environ.get(BACKEND_ENV_VAR, DEFAULT_BACKEND_NAME)
+    if choice not in VALID_BACKEND_NAMES:
+        raise ValueError(
+            "Invalid backend {!r}. Expected one of {}. "
+            "Set the constructor kwarg _backend=, or the {} environment variable.".format(
+                choice, VALID_BACKEND_NAMES, BACKEND_ENV_VAR
+            )
+        )
+    return choice
+
+
+def make_backend(explicit: Optional[str]) -> CosmosBackend:
+    """Build the single backend instance a sync ``CosmosClient`` will hold.
+
+    ``explicit`` is what the caller passed as ``_backend=`` (or ``None`` if
+    they passed nothing — in which case the env var, then the default, are
+    consulted). The returned object is either a ``CorePythonBackend`` or a
+    ``RustBackend``.
+    """
+    name = resolve_backend_name(explicit)
+    if name == BACKEND_NAME_RUST:
+        return RustBackend()
+    return CorePythonBackend()
+

sdk/cosmos/azure-cosmos/azure/cosmos/_backend/rust.py

@@ -0,0 +1,62 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+"""Sync Rust backend.
+
+This is the only Python module allowed to import the compiled PyO3
+module ``azure.cosmos._azure_cosmos_pyo3``. The import-guard test
+(``tests/test_backend_wiring_unit.py``) enforces that rule by walking
+every ``.py`` file under ``azure/cosmos/`` and failing the build if any
+other module reaches across the abstraction.
+
+The compiled PyO3 module may not be present in every checkout (a fresh
+clone will not have run ``maturin develop`` yet). The import below is
+guarded with ``try / except ImportError`` so this file can still be
+loaded; ``RustBackend.create_item`` will then raise
+``NotImplementedError`` for any caller that asks for the Rust backend.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+from .base import BackendResponse, CosmosBackend, PreparedRequest
+from .constants import BACKEND_NAME_RUST
+
+_LOGGER = logging.getLogger(__name__)
+
+# Module-level reference set once at import time, under the GIL. Read-only
+# afterwards, so it is safe to share across threads and across clients.
+_pyo3_driver: Optional[Any] = None
+try:
+    from azure.cosmos import _azure_cosmos_pyo3  # type: ignore[attr-defined]
+    _pyo3_driver = _azure_cosmos_pyo3
+except ImportError:
+    _LOGGER.debug(
+        "_azure_cosmos_pyo3 module not available; RustBackend operations "
+        "will raise NotImplementedError until the PyO3 wrapper is built."
+    )
+
+
+class RustBackend(CosmosBackend):
+    """Routes ``create_item`` calls through the in-tree Rust driver.
+
+    Today this class's ``create_item`` raises ``NotImplementedError`` on
+    every call. That is the expected behavior for the dispatch-only
+    slice — a developer who runs the existing test suite with
+    ``COSMOS_BACKEND=rust`` should see every create_item test fail
+    loudly, which proves the dispatch wiring works end-to-end. The real
+    implementation lands once the Rust-side gaps are closed and the
+    helper layer exists.
+    """
+
+    name = BACKEND_NAME_RUST
+
+    def __init__(self) -> None:
+        self._driver = _pyo3_driver
+
+    def create_item(self, prepared: Optional[PreparedRequest]) -> Optional[BackendResponse]:
+        raise NotImplementedError("RustBackend.create_item not implemented yet")
+