fix: pass context to describe_collection in SearchIteratorV2 (#3271)

.gitignore

@@ -25,6 +25,8 @@ venv/
 
 # Local Temp
 temp/
+.worktrees/
+.worktree/
 *.swp
 assets/
 TODO

pymilvus/bulk_writer/constants.py

@@ -56,13 +56,13 @@
     DataType.JSON.name: lambda x: isinstance(x, (str, list, dict)),
     DataType.TIMESTAMPTZ.name: lambda x: isinstance(x, str),
     DataType.GEOMETRY.name: lambda x: isinstance(x, str),
-    DataType.STRUCT.name: lambda x, max_cap: struct_validator(x, max_cap),
-    DataType.FLOAT_VECTOR.name: lambda x, dim: float_vector_validator(x, dim),
-    DataType.BINARY_VECTOR.name: lambda x, dim: binary_vector_validator(x, dim),
+    DataType.STRUCT.name: struct_validator,
+    DataType.FLOAT_VECTOR.name: float_vector_validator,
+    DataType.BINARY_VECTOR.name: binary_vector_validator,
     DataType.FLOAT16_VECTOR.name: lambda x, dim: float16_vector_validator(x, dim, False),
     DataType.BFLOAT16_VECTOR.name: lambda x, dim: float16_vector_validator(x, dim, True),
-    DataType.SPARSE_FLOAT_VECTOR.name: lambda x: sparse_vector_validator(x),
-    DataType.INT8_VECTOR.name: lambda x, dim: int8_vector_validator(x, dim),
+    DataType.SPARSE_FLOAT_VECTOR.name: sparse_vector_validator,
+    DataType.INT8_VECTOR.name: int8_vector_validator,
     DataType.ARRAY.name: lambda x, cap: (isinstance(x, (list, np.ndarray)) and len(x) <= cap),
 }
 

pymilvus/client/search_iterator.py

@@ -49,7 +49,7 @@ def __init__(
             self._left_res_cnt = limit
 
         self._conn = connection
-        self._set_up_collection_id(collection_name)
+        self._set_up_collection_id(collection_name, **kwargs)
         kwargs[COLLECTION_ID] = self._collection_id
         self._params = {
             "collection_name": collection_name,
@@ -73,8 +73,8 @@ def __init__(
         self._batch_size = batch_size
         self._probe_for_compability(self._params)
 
-    def _set_up_collection_id(self, collection_name: str):
-        res = self._conn.describe_collection(collection_name)
+    def _set_up_collection_id(self, collection_name: str, **kwargs):
+        res = self._conn.describe_collection(collection_name, **kwargs)
         self._collection_id = res[COLLECTION_ID]
 
     def _check_token_exists(self, token: Union[str, None]):

tests/test_search_iterator.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from pymilvus.client.call_context import CallContext
 from pymilvus.client.search_iterator import SearchIteratorV2
 from pymilvus.client.search_result import SearchResult
 from pymilvus.exceptions import ParamError, ServerVersionIncompatibleException
@@ -88,6 +89,20 @@ def test_multiple_vectors_error(self, mock_connection):
                 batch_size=100,
             )
 
+    def test_context_passed_to_describe_collection(self, mock_connection, search_data):
+        """Regression test for #3270: context (db_name) must be forwarded to describe_collection"""
+        ctx = CallContext(db_name="nondefault")
+
+        SearchIteratorV2(
+            connection=mock_connection,
+            collection_name="test_collection",
+            data=search_data,
+            batch_size=100,
+            context=ctx,
+        )
+
+        mock_connection.describe_collection.assert_called_once_with("test_collection", context=ctx)
+
     @patch("pymilvus.client.search_iterator.SearchIteratorV2._probe_for_compability")
     def test_next_without_external_filter(self, mock_probe, mock_connection, search_data):
         mock_connection.search.return_value = self.create_mock_search_result()