From 1dd80d57332b279a0fa7bdaafd4f3be27728c1dd Mon Sep 17 00:00:00 2001 From: bharath <9440110838bharath@gmail.com> Date: Fri, 17 Jul 2026 23:50:11 -0400 Subject: [PATCH 1/2] Reuse Milvus client in vector sink write path _MilvusSink.write() constructed a new MilvusClient on every call, which discarded the retry-wrapped client created in __enter__. This bypassed the connection retry/backoff and left the initial client unclosed (leaked), while __exit__ only closed the last client. The write() docstring also incorrectly claimed the collection was flushed for durability. Reuse the client established in __enter__ for the upsert, and correct the docstring to describe the actual (no-flush) visibility semantics. Add a unit test asserting write() reuses the __enter__ client instead of reconstructing one per call. Co-authored-by: Cursor --- .../ml/rag/ingestion/milvus_search.py | 12 ++++--- .../ml/rag/ingestion/milvus_search_test.py | 34 +++++++++++++++++++ 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py index b7cad3796b13..0516e0525f96 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py @@ -287,16 +287,18 @@ def write(self, documents): """Writes a batch of documents to the Milvus collection. Performs an upsert operation to insert new documents or update existing - ones based on primary key. After the upsert, flushes the collection to - ensure data persistence. + ones based on primary key, reusing the client established in ``__enter__``. + + Note: + This submits the upsert to Milvus; it does not call ``flush()`` on the + collection. Newly written rows become queryable according to Milvus' + normal segment-sync behavior. Callers that require immediate + read-after-write visibility should flush the collection explicitly. Args: documents: List of dictionaries representing Milvus records to write. Each dictionary should contain fields matching the collection schema. """ - self._client = MilvusClient( - **unpack_dataclass_with_kwargs(self._connection_params)) - resp = self._client.upsert( collection_name=self._write_config.collection_name, partition_name=self._write_config.partition_name, diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py index 80d55ac9382c..e313118dac13 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py @@ -15,12 +15,15 @@ # limitations under the License. # import unittest +from unittest import mock from parameterized import parameterized try: + from apache_beam.ml.rag.ingestion import milvus_search as milvus_search_module from apache_beam.ml.rag.ingestion.milvus_search import MilvusVectorWriterConfig from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig + from apache_beam.ml.rag.ingestion.milvus_search import _MilvusSink from apache_beam.ml.rag.utils import MilvusConnectionParameters except ImportError as e: raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}') @@ -119,5 +122,36 @@ def test_invalid_configuration_parameters( self.assertIn(expected_error_msg, str(context.exception)) +class TestMilvusSinkClientReuse(unittest.TestCase): + """Unit tests for Milvus sink client reuse. + + Verifies that ``_MilvusSink.write`` reuses the client established in + ``__enter__`` instead of constructing (and leaking) a new client on every + write. + """ + def _sink(self): + connection_params = MilvusConnectionParameters(uri="http://localhost:19530") + write_config = MilvusWriteConfig(collection_name="test_collection") + return _MilvusSink(connection_params, write_config) + + def test_write_reuses_enter_client(self): + """write() must not construct a new MilvusClient per call.""" + with mock.patch.object(milvus_search_module, + 'MilvusClient') as mock_client_cls: + with self._sink() as sink: + # One client created on __enter__. + self.assertEqual(mock_client_cls.call_count, 1) + + sink.write([{"id": 1}]) + sink.write([{"id": 2}]) + + # write() reuses that client; no new construction. + self.assertEqual(mock_client_cls.call_count, 1) + self.assertEqual(mock_client_cls.return_value.upsert.call_count, 2) + + # Client closed once on __exit__. + self.assertEqual(mock_client_cls.return_value.close.call_count, 1) + + if __name__ == '__main__': unittest.main() From 534ba903e7b7f8104e54a1aa56e9d67b0e0e23c7 Mon Sep 17 00:00:00 2001 From: bharath <9440110838bharath@gmail.com> Date: Sat, 18 Jul 2026 23:52:58 -0400 Subject: [PATCH 2/2] Clear Milvus client after close in sink __exit__ Set self._client to None after closing so a reused sink instance does not keep a stale closed client. Address review feedback. Co-authored-by: Cursor --- sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py | 1 + sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py index 0516e0525f96..cfca29a1ada6 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py @@ -348,3 +348,4 @@ def __exit__(self, exc_type, exc_val, exc_tb): _ = exc_type, exc_val, exc_tb # Unused parameters if self._client: self._client.close() + self._client = None diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py index e313118dac13..0c67e15823eb 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py @@ -149,8 +149,9 @@ def test_write_reuses_enter_client(self): self.assertEqual(mock_client_cls.call_count, 1) self.assertEqual(mock_client_cls.return_value.upsert.call_count, 2) - # Client closed once on __exit__. + # Client closed once on __exit__, and cleared so the sink can be reused. self.assertEqual(mock_client_cls.return_value.close.call_count, 1) + self.assertIsNone(sink._client) if __name__ == '__main__':