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..cfca29a1ada6 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, @@ -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 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..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 @@ -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()