[Python] Reuse Milvus client in RAG vector sink write path (#39378)

* 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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: bharath000

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,
@@ -346,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

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,37 @@ 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__, 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__':
   unittest.main()