python-rust-integration - #46851
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces an internal “backend dispatch” layer for azure-cosmos (sync + async) to allow selecting between the existing core-Python implementation and an opt-in Rust/PyO3-backed path, including per-request backend tagging in the User-Agent and a unit-test suite to pin down the wiring/selection rules.
Changes:
- Add sync/async backend abstractions, constants, and factories with selection precedence
_backendkwarg >COSMOS_BACKENDenv var > defaultcore-python. - Wire backend selection into
CosmosClient/async client construction, expose backend instances onclient_connection, and addcreate_itemdispatch + per-request backend stamping. - Extend
CosmosUserAgentPolicyto appendbackend=<name>and add fast unit tests (including an import-guard) to verify wiring and behavior.
Reviewed changes
Copilot reviewed 15 out of 17 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/cosmos/azure-cosmos/tests/test_backend_wiring_unit.py | New unit tests covering import-guard, backend selection precedence/validation, container dispatch behavior, and UA stamping. |
| sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py | Selects/initializes sync backend(s) at client construction and attaches them to client_connection. |
| sdk/cosmos/azure-cosmos/azure/cosmos/container.py | Adds sync create_item dispatch and stamps the chosen backend into per-request options. |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py | Selects/initializes async backend(s) at client construction and attaches them to client_connection. |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | Adds async create_item dispatch and stamps the chosen backend into per-request options. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_user_agent_policy.py | Appends a backend=<name> token to the User-Agent based on per-request options. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/base.py | Defines the sync backend ABC and shared immutable request/response dataclasses. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/constants.py | Centralizes backend names, env var, and per-request option key. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/factory.py | Implements backend name resolution + sync backend instance factory. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/core_python.py | Sync “core-python” backend stub (returns None to fall through to existing path). |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/rust.py | Sync Rust backend stub with guarded PyO3 import and loud NotImplementedError. |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_backend/base.py | Defines the async backend ABC and re-exports shared dataclasses. |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_backend/factory.py | Async backend factory mapping resolved backend name to async backend instance. |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_backend/core_python.py | Async “core-python” backend stub (returns None to fall through). |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_backend/rust.py | Async Rust backend stub with guarded PyO3 import and loud NotImplementedError. |
| sdk/cosmos/azure-cosmos/azure/cosmos/_backend/init.py | Backend package marker/export surface (no diff shown). |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_backend/init.py | Async backend package marker/export surface (no diff shown). |
| # Pick which backends this client will hold. Two attributes, named | ||
| # after the two backend types: | ||
| # - _core_python_backend: always present. The default, and the | ||
| # always-available path for any request whose kwargs the Rust | ||
| # backend doesn't support yet. | ||
| # - _rust_backend: present only when Rust is selected as default. | ||
| # Its presence (vs None) is the "Rust is default" signal — no | ||
| # extra attribute needed. | ||
| # Precedence for selection: kwarg `_backend=` > COSMOS_BACKEND env | ||
| # var > "core-python". | ||
| backend_choice = kwargs.pop("_backend", None) |
| self._rust_backend: Optional[RustBackend] = ( | ||
| chosen if isinstance(chosen, RustBackend) else None | ||
| ) | ||
| logging.getLogger(__name__).info( |
| self._rust_backend: Optional[AsyncRustBackend] = ( | ||
| chosen if isinstance(chosen, AsyncRustBackend) else None | ||
| ) | ||
| logging.getLogger(__name__).info( |
| backend = core_python_backend | ||
| if backend is not None: | ||
| # Stamp the backend that actually handled this call so the | ||
| # user-agent policy can append `; backend=<name>` per request. |
| backend = core_python_backend | ||
| if backend is not None: | ||
| # Stamp the backend that actually handled this call so the | ||
| # user-agent policy can append `; backend=<name>` per request. |
| # Compose any per-request user-agent suffixes. Two sources today: | ||
| # 1. Cosmos feature flags derived from the global endpoint manager | ||
| # (existing behavior — circuit breaker, PPAF, etc.). | ||
| # 2. The backend that handled this request ("core-python" or "rust"), | ||
| # stamped per-request from the dispatch site so server-side | ||
| # logs reflect the path the call actually took even when a | ||
| # Rust-default client falls back to core-python for a single call. | ||
| suffix_parts = [] |
| with pytest.raises(Exception): # FrozenInstanceError | ||
| setattr(p, "body_bytes", b"different") | ||
|
|
||
| r = BackendResponse(status_code=201) | ||
| with pytest.raises(Exception): | ||
| setattr(r, "status_code", 200) |
| > declaration must be changed to the approved published `azure_data_cosmos_driver` version; | ||
| > the sibling path remains a local-development arrangement only. | ||
|
|
||
| **The problem:** these are *two separate crates*, and the binding depends on the driver. |
There was a problem hiding this comment.
For Development: the binding crate's Cargo.toml points to a released version of the driver on crates.io or github+ref for the version of a live development version of the driver on azure-sdk-for-rust
For release: Only point to a released version on crates.io. I would also think that a GA version of the Python wheel and binding crate would only depend on GA versions of the driver (that is, a GA version of the binding crate CANNOT depend on a beta/preview version of the driver)
|
|
||
| > **How the current build resolves the driver.** The binding's `Cargo.toml` declares | ||
| > `azure_data_cosmos_driver` as a local **path dependency** pointing to the sibling | ||
| > `azure-sdk-for-rust` checkout. Maturin starts the binding build, then Cargo compiles both |
There was a problem hiding this comment.
Current binary builds for Python Storage Extension (native C) use cibuildwheel ... as long as Maturin is supported there, the change here will probably look the same as the changes we've made to ship Python Storage Extension.
|
|
||
| | File | What it declares for `tokio` | | ||
| |---|---| | ||
| | `azure_cosmos_rust/Cargo.toml` (the binding crate) | *which* crate to depend on, and which features it needs — `rt-multi-thread`, `macros` | |
There was a problem hiding this comment.
Should this path be sdk/cosmos/azure-cosmos/azure_cosmos_rust/Cargo.toml
| **The driver revision is also not pinned by Cargo.** The dependency is a relative `path`, not | ||
| a `git` dependency with a `rev`. The Python commit therefore records neither | ||
| `5c170b538` nor any other Rust commit; it compiles whichever sibling checkout happens to be | ||
| at that path. The release fix is to replace this path with the approved crates.io version. | ||
| If developers still want to test unreleased driver source, they can override that dependency | ||
| locally without putting the override in the published manifest. |
There was a problem hiding this comment.
This is reasonable for local development but will fail if you try to check it into git and build in a pipeline. See comment above about taking a dependency on crates or from source in github.
| binding still compiles against the version selected for release. | ||
|
|
||
| **Verify resolution on a clean build, don't assume it.** The branch does not commit | ||
| **`Cargo.lock`**, the generated file recording every crate in the graph at one exact version. |
There was a problem hiding this comment.
Work with Heath Stewart (@heaths) and Kashif Khan (@kashifkhan) about locking dependencies for build and release and how the movement of those should be expressed in files in the repo around release time.
There was a problem hiding this comment.
Cargo.lock should be checked in for both bin and lib crates. That guidance changed to include lib crates years ago and, really, if you're producing wheels you are shipping bins anyway, per my limited understanding of wheels.
There was a problem hiding this comment.
Just make sure the dependabot.yml configuration is set up to review and update any Cargo.toml/Cargo.lock files.
| emphasize that it is machine code built for one OS/CPU; "extension" is Python's word for a | ||
| module written in compiled code rather than `.py`. | ||
|
|
||
| So Cargo's `azure_cosmos_rust.dll` has to be **renamed** to `_rust.pyd` and **placed** |
There was a problem hiding this comment.
Worth considering if the build system itself can be configured to produce this file so you aren't writing extra tooling to discover and rename files in the pipeline engsys.
| Cargo alone does not perform all three: | ||
|
|
||
| 1. Run Cargo to compile the `cdylib`. | ||
| 2. Rename the output (`azure_cosmos_rust.dll` → `_rust.pyd`) and copy it into `azure/cosmos/`. |
There was a problem hiding this comment.
Hopefully this can be done as part of the build system using Maturin, crate, etc.
| | Windows 64-bit | `win_amd64` | | ||
| | Linux x86_64 | `manylinux_2_17_x86_64` | | ||
| | Linux ARM64 | `manylinux_2_17_aarch64` | | ||
| | macOS Intel | `macosx_10_12_x86_64` | |
There was a problem hiding this comment.
Kashif Khan (@kashifkhan) -- is this still a target for Python packages?
There was a problem hiding this comment.
no, macos is ARM only
| | macOS Intel | `macosx_10_12_x86_64` | | ||
| | macOS Apple Silicon | `macosx_11_0_arm64` | | ||
|
|
||
| The exact platform tags and supported architectures must be confirmed with the release |
There was a problem hiding this comment.
What statements are being made about supported Python versions, OS, and hardware platforms? (CPython versions, ABI, PyPy, etc.)
There was a problem hiding this comment.
added explicitly to the document
| same way. | ||
| 2. **Where will the sdist get the Rust driver source?** This is now decided: release builds | ||
| will depend on the **published crates.io driver** rather than the local sibling path. The | ||
| binding's `Cargo.toml` must replace `path = "../../../../../azure-sdk-for-rust/..."` with |
There was a problem hiding this comment.
This should probably be shipped from main with the correct source code checked in and not referencing something on a dev machine.
There was a problem hiding this comment.
azure_data_cosmos_driver = {
version = "1.2.3",
features = ["__internal_native_query_plan"]
}yes it would be something like this local folder is used for dev purposes for faster iterations
|
|
||
| ## 11. The release math changes: today vs. after v5 | ||
|
|
||
| The release shape changes enough that the pure-Python assumptions ("one wheel, one active |
There was a problem hiding this comment.
Much of this work is already done today for packages that are detected as building platform-specific wheels (for example, azure-storage-extensions). Therefore, it's easiest if these wheels can be produced by existing processes supported by the build system.
| 1. Select the native multi-platform build path for `azure-cosmos`. | ||
| 2. Make the package parser/builder recognize the Maturin extension and invoke its PEP | ||
| 517/Maturin build through `cibuildwheel`. | ||
| 3. Install or provide Rust. |
There was a problem hiding this comment.
Which version of rust should be used to build? MSRV? Stable? etc.
| 3. Install or provide Rust. | ||
| 4. Build against the approved published crates.io driver version; remove any dependency on a | ||
| sibling `azure-sdk-for-rust` checkout from the release package and CI job. | ||
| 5. Build and test every approved platform wheel. |
There was a problem hiding this comment.
Build machines are:
| OS | Architecture |
|---|---|
| windows | x86_64 |
| linux | x86_64 |
| macos | ARM64 |
Only those platforms can be directly tested. Platforms like ARM64 Windows and Linux are not currently supported in the CI system.
| 4. Build against the approved published crates.io driver version; remove any dependency on a | ||
| sibling `azure-sdk-for-rust` checkout from the release package and CI job. | ||
| 5. Build and test every approved platform wheel. | ||
| 6. Aggregate and publish only after all required artifacts pass validation. |
There was a problem hiding this comment.
Step 7: Unpack, sign native extensions, repack wheels (Windows, macos)
|
|
||
| `cibuildwheel` is already present in the central CI tool set, but Cosmos has no package-level | ||
| configuration for it. Maturin is declared as a build-system requirement in the package | ||
| `pyproject.toml`; whether EngSys also wants it preinstalled on agents is an implementation |
There was a problem hiding this comment.
Can the cibuildwheel config and matrix be added to pyproject.toml?
| - Without `ENABLE_EXTENSION_BUILD`, only Linux runs package generation; Windows and macOS | ||
| publish empty package artifacts. | ||
| - With only that variable enabled, all three jobs run `sdk_build`, but the current package | ||
| parser still sees `setup.py` with zero `ext_modules`, so it does not select | ||
| `cibuildwheel`. | ||
| - After both gates are fixed, each job can contribute its approved wheel or wheels to the | ||
| existing aggregation path. Repository source proves CI artifact collation, but it does not | ||
| by itself prove production signing or a complete multi-file PyPI release. | ||
|
|
||
| The intended end state is therefore not a new aggregation design. It is a correctly detected | ||
| Maturin package feeding validated native wheels into the aggregation design that already | ||
| exists. |
There was a problem hiding this comment.
The build system will be adjusted to support landing appropriate toolchains and building (ideally with cibuildwheel as it's already supported, assuming it can orchestrate the output matrix)
|
|
||
| 1. **Agent toolchain.** Do the Windows, Linux, and macOS agents provide a Rust toolchain at | ||
| or above the minimum the *driver* requires (§3 — higher than the version our own binding | ||
| workspace currently declares)? If not, should each build install it, or should the agent |
There was a problem hiding this comment.
Once a version is specified, it can be installed. This should probably be done with something like a rust-toolchain.toml file checked in either at the level of the package or near the repo root.
| 2. **ARM coverage — ARM and x86 are two different CPU designs, and machine code built for | ||
| one will not run on the other. So who builds the ARM wheels? Either we rent an ARM machine | ||
| and build there (simplest, needs that agent to exist), or we build on our existing x86 | ||
| machines by telling the compiler "target ARM" (cross-compilation — no new machine, but | ||
| nothing on that box can actually *run* the result to test it), or we fake an ARM CPU in | ||
| software (emulation — it runs, but slowly). |
There was a problem hiding this comment.
macos -- building on ARM64 platform, no problem
linux -- cibuildwheel orchestrates containers and emulators for an emulated native build
windows -- cibuildwheel orchestrates a cross-compile with toolchains on the agent machine
This is a big part of why I recommend using cibuildwheel
| 3. **Signing and multi-file publication.** The CI templates prove that platform artifacts | ||
| can be collected. Confirm separately that release tooling signs the native binaries as | ||
| required, publishes the complete required artifact set as one version, and fails safely | ||
| if any required wheel is missing. |
There was a problem hiding this comment.
Signing is currently working in a branch, we'll have to do some re-working to support a mixture of bdist and sdist in various services.
|
|
||
| --- | ||
|
|
||
| ## 16. Rolling it out without breaking a release |
There was a problem hiding this comment.
This is for EngSys to handle... It amounts to:
- Add support for detecting rust packaging for azure-cosmos and call cibuildwheel (preferably) to generate relevant packages
- Add support for signing artifacts inside bdist wheels
- Support publishing a mixture of signed bdist wheels and sdist wheels
migrating from core-python to rust native