Skip to content

enhance: support idempotent requests for DML - #3784

Open
bigsheeper wants to merge 1 commit into
milvus-io:masterfrom
bigsheeper:feat/idempotency-key
Open

enhance: support idempotent requests for DML#3784
bigsheeper wants to merge 1 commit into
milvus-io:masterfrom
bigsheeper:feat/idempotency-key

Conversation

@bigsheeper

Copy link
Copy Markdown
Contributor

Milvus accepts a client idempotency key as transport metadata: the REST header Idempotency-Key and the gRPC metadata key idempotency-key (milvus-io/milvus#52544). A retry carrying the same key resolves to the original request instead of doing the work again. This PR exposes it from pymilvus.

Changes

  • MilvusClient / AsyncMilvusClient insert, upsert and delete accept idempotency_key as a keyword argument, following the client_request_id precedent: CallContext carries it and emits the idempotency-key gRPC metadata only when set, so a keyless call is byte-identical to today.
  • The ORM path picks it up through connections._generate_call_context, so Collection.insert(..., idempotency_key=...) works too.
  • bulk_import(idempotency_key=...) sends the Idempotency-Key header.
client.insert("events", rows, idempotency_key="nightly-2026-09-07-batch-3")

bulk_import(url=url, collection_name="events", files=files,
            idempotency_key="nightly-2026-09-07-batch-3")

Server-side status

Tests

  • tests/unit/test_idempotency_key.py (new): CallContext metadata emission, sync and async client insert/upsert/delete forwarding, ORM context generation.
  • tests/unit/test_bulk_import.py: header emission at _http_headers, _post_request and bulk_import.
  • Full tests/unit run: 4718 passed; ruff and black clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01WQihG8uqLQb8LWW5aVivjY

Milvus accepts a client idempotency key as transport metadata: the REST
header Idempotency-Key and the gRPC metadata key idempotency-key
(milvus-io/milvus#52544). A retry carrying the same key resolves to the
original request instead of doing the work again. This exposes it from
pymilvus.

- MilvusClient / AsyncMilvusClient insert, upsert and delete accept
  idempotency_key as a keyword argument, following the client_request_id
  precedent: CallContext carries it and emits the idempotency-key gRPC
  metadata only when set, so a keyless call is byte-identical to today.
- The ORM path picks it up through connections._generate_call_context.
- bulk_import(idempotency_key=...) sends the Idempotency-Key header.

Bulk import honors the key on merged servers. Insert honors it once
milvus-io/milvus#50007 lands. Delete and upsert carry the key but the
proxy does not act on it yet.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WQihG8uqLQb8LWW5aVivjY
Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
@sre-ci-robot

Copy link
Copy Markdown

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: bigsheeper
To complete the pull request process, please assign xuanyang-cn after the PR has been reviewed.
You can assign the PR to them by writing /assign @xuanyang-cn in a comment when ready.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@mergify

mergify Bot commented Sep 7, 2026

Copy link
Copy Markdown

Tick the box to add this pull request to the merge queue (same as @mergifyio queue).

  • Queue this pull request

@codecov

codecov Bot commented Sep 7, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.22%. Comparing base (a4b3b38) to head (f3ee350).

Additional details and impacted files
@@           Coverage Diff           @@
##           master    #3784   +/-   ##
=======================================
  Coverage   94.22%   94.22%           
=======================================
  Files          77       77           
  Lines       16204    16211    +7     
=======================================
+ Hits        15268    15275    +7     
  Misses        936      936           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@mergify mergify Bot added the ci-passed label Sep 7, 2026
@czs007

czs007 commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Must-fix issues introduced by this PR: 0
Merge recommendation: Mergeable as-is — this review found no must-fix issue introduced by this PR.

Adversarial review found no issues requiring changes.

Verified:

  • pymilvus/client/call_context.py:19-27 — the idempotency-key header is appended only when the key is non-empty, so keyless calls emit exactly the same three metadata tuples as before this PR.
  • pymilvus/orm/connections.py:592-596kwargs.get("idempotency_key", "") mirrors the client_request_id precedent; the key stays in **kwargs and every downstream terminal (Prepare.delete_request, describe_collection, MutationFuture) absorbs it, so no TypeError is possible.
  • pymilvus/client/grpc_handler.py:972-976insert_rows (and likewise upsert_rows at :1315, batch_insert at :1101, delete at :1149) issues a single RPC per call, so no multiple chunks ever share one key.
  • pymilvus/milvus_client/milvus_client.py:885-907 — delete-by-ids builds the call context twice but both carry the same key and only one Delete RPC is sent.
  • pymilvus/client/bulk_import.py:76-81_post_request consumes the key via pop into the HTTP header only; tests/unit/test_bulk_import.py:216-217 asserts it never reaches the JSON body or lingers in kwargs.
  • tests/unit/test_idempotency_key.py_context_of reads the real context= keyword argument, matching the call shape at milvus_client.py:280, so the sync/async insert/upsert/delete assertions exercise the actual forwarding path.

Low — non-blocking

  • pymilvus/client/grpc_handler.py:1041 — The same CallContext that carries idempotency-key is passed into schema resolution (insert_rows_prepare_row_insert_request_get_schema at :1000-1002, upsert_rows at :1280-1282, ORM columnar batch_insert at :1071-1073, MilvusClient.delete(ids=…) at milvus_client.py:884-889, and the async handler at async_grpc_handler.py:707-709). On a schema cache miss, describe_collection sends _api_level_md(context) at :820-822, so DescribeCollection goes out with the same idempotency-key as the following Insert/Upsert/Delete. This contradicts the PR description's statement that the key is gRPC metadata for that DML only. Server-side dedup is keyed by (messageType, scopedKey) and only broadcast write requests enter the index, and the proxy DescribeCollection path never reads this metadata, so no incorrect behavior is reproducible today; this is a contract/consistency gap rather than a defect. Suggestion: derive a key-less CallContext for the schema-resolution path (or strip the field inside _get_schema), and add a negative assertion that DescribeCollection metadata contains no idempotency-key on a cache miss. (raised by xiaocai2333)
  • tests/unit/test_idempotency_key.py:140 — The ORM entry points are only covered at the connections._generate_call_context layer (:140-143); nothing asserts that Collection.insert / upsert / delete_get_connection(**kwargs) (collection.py:557/635/695) → _generate_call_context (:194) → handler actually delivers a context carrying the key. The chain is correct on reading, but collection.py already has several _get_connection() calls without kwargs (:263/:309/:1219/:1361), and a future change of that kind on the DML methods would go uncaught. Suggestion: mirror TestMilvusClient with a mock handler and assert the context passed to insert_rows / batch_insert / upsert_rows / delete contains the key. (raised by czs007, xiaocai2333)
  • tests/unit/test_idempotency_key.py:50 — All assertions are positive checks on the DML context, and both the sync fixture (:50) and the async fixture (:96) mock _get_schema outright, so describe_collection is never reached and the leak above is invisible to the suite. Suggestion: add one test that uses a real GrpcHandler (or mocks only _stub), clears GlobalCache.schema to force a cache miss, and asserts _stub.DescribeCollection metadata lacks idempotency-key while _stub.Insert metadata carries it. That test fails on the current code and pairs with the first item. (raised by czs007, xiaocai2333)

@tinswzy

tinswzy commented Sep 9, 2026

Copy link
Copy Markdown

Reviewed against the server side of milvus-io/milvus#50007 and #52544. The plumbing is
right, and one thing that matters most is right: context is built once in the client
method and @retry_on_rpc_failure() sits on insert_rows, so every retry re-emits the
same key. That is the property the whole feature rests on.

The wire name matches the server constants exactly, and empty-vs-absent is handled the way
the server reads it. What follows is about what the parameter promises versus what the
server does with it.

Must fix

1. insert with a key is rejected on a default-configured server — it is not ignored.

// task_insert.go:192 -> task_insert_idempotency.go:65
if it.idempotencyKey != "" && !globalIdempotencyEnabled {
    return merr.WrapErrParameterInvalidMsg("idempotency key is not accepted when idempotent write is globally disabled")
}
if it.idempotencyKey != "" && !collectionIdempotencyEnabled {
    return merr.WrapErrParameterInvalidMsg("idempotency key is not accepted when collection idempotent write is disabled")
}

Both gates default off: streaming.idempotency.enabled is false, and the collection
property is unset. So adding this argument to working code makes the write start failing
with error 1100. The docstring should say both switches are required and that sending a key
early is an error, not a no-op. (milvus-docs#3635 already states this — worth matching.)

2. upsert and delete docstrings should say the server does not act on the key.

GetIdempotencyKeyFromContext has exactly one caller, Proxy.Insert. Upsert and delete
never read it. Today the same parameter has three behaviours — insert errors when disabled,
upsert and delete silently drop it — behind three docstrings that read alike. Someone who
sets a key on delete and retries gets a second delete with no signal.

Worth fixing

3. Nothing validates the key before it reaches the server. ValidateIdempotencyKey
requires printable ASCII (0x20–0x7E) and a length bound (streaming.idempotency.maxKeyLength,
256 by default). A Python caller passing a non-ASCII string — a UUID is fine, a batch name
with a non-Latin character is not — gets a 1100 from deep in the server. A few lines client
side would give a better error.

4. idempotency_key is kwargs-only, so a misspelling silently disables idempotency.
idempotencyKey=, idempotency_id=, or **{"idempotency-key": ...} are all absorbed by
**kwargs at both ends — _generate_call_context does .get("idempotency_key", "") and
insert_rows(**kwargs) swallows the rest. No error, and the write is simply not idempotent.

client_request_id has the same shape, but the consequences differ: a lost request id costs
a log correlation, a lost idempotency key costs duplicate rows. Making it an explicit keyword
parameter would not stop **kwargs from swallowing a typo — that is inherent here — but it
would put the right name in autocomplete, help() and the signature, which is where the
name is most likely to be gotten right.

Relatedly: _generate_call_context accepts both client_request_id and client-request-id
but only the underscore form for the new parameter. Rather than accepting both, consider
raising on idempotency-key in kwargs — readers meet that spelling as the wire name and
copying it currently fails silently.

5. The Go SDK and this PR disagree on upsert. client/milvusclient/write_options.go
rejects a key on Upsert client-side ("idempotency key is only supported for Insert"), while
this sends it and the server drops it. Two official SDKs, opposite contracts for one
operation. Silently dropping is the weakest of the three options; worth a maintainer call
rather than settling it in either PR alone.

Optional

6. _generate_call_context(**kwargs) appears 116 times in the sync client and 118 in the
async one, so search(..., idempotency_key=...) and create_collection(...) also emit the
metadata — undocumented and untested. On the DDL path #52544's broadcaster does act on a key,
so this is either a feature or a surprise; worth deciding deliberately.

7. Only async upsert got the docstring note; async insert and delete did not, though
the description covers all three.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants