Skip to content

[Bug]: Milvus RAG vector sink recreates client on every write, bypassing retry and leaking the connection #39377

Description

@bharath000

What happened?

In the Milvus RAG vector sink (apache_beam/ml/rag/ingestion/milvus_search.py), _MilvusSink.write() creates a new MilvusClient on every call:

def write(self, documents):
    self._client = MilvusClient(
        **unpack_dataclass_with_kwargs(self._connection_params))
    resp = self._client.upsert(...)

But __enter__ already creates the client using retry_with_backoff. So the write path:

  • throws away the retry-wrapped client from __enter__ and makes a fresh one without retry
  • leaves the original client unclosed, so it leaks (MilvusClient has no __del__)
  • __exit__ only closes the last client

The write() docstring also says it "flushes the collection to ensure data persistence", but there is no flush() call anywhere.

Expected: write() should reuse the client from __enter__, and the docstring should describe the real behavior (no flush).

Issue Priority

Priority: 3 (minor)

Issue Components

  • Component: Python SDK
  • Component: Java SDK
  • Component: Go SDK
  • Component: Typescript SDK
  • Component: IO connector
  • Component: Beam YAML
  • Component: Beam examples
  • Component: Beam playground
  • Component: Beam katas
  • Component: Website
  • Component: Infrastructure
  • Component: Spark Runner
  • Component: Flink Runner
  • Component: Prism Runner
  • Component: Twister2 Runner
  • Component: Hazelcast Jet Runner
  • Component: Google Cloud Dataflow Runner

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions