locust 2.38.1.dev4__py3-none-any.whl → 2.38.2.dev7__py3-none-any.whl

locust/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '2.38.1.dev4'
-__version_tuple__ = version_tuple = (2, 38, 1, 'dev4')
+__version__ = version = '2.38.2.dev7'
+__version_tuple__ = version_tuple = (2, 38, 2, 'dev7')
+
+__commit_id__ = commit_id = None

locust/contrib/fasthttp.py

@@ -231,7 +231,7 @@ class FastHttpSession:
         elif self.auth_header:
             headers["Authorization"] = self.auth_header
         if "Accept-Encoding" not in headers and "accept-encoding" not in headers:
-            headers["Accept-Encoding"] = "gzip, deflate, br, zstd"
+            headers["Accept-Encoding"] = "gzip, deflate, br"
 
         if not data and json is not None:
             data = unshadowed_json.dumps(json)

locust/contrib/milvus.py

@@ -0,0 +1,407 @@
+import gevent.monkey
+
+gevent.monkey.patch_all()
+import grpc.experimental.gevent as grpc_gevent
+
+grpc_gevent.init_gevent()
+
+from locust import User, events
+
+import time
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pymilvus import CollectionSchema, MilvusClient
+from pymilvus.milvus_client import IndexParams
+
+
+class BaseClient(ABC):
+    @abstractmethod
+    def close(self) -> None:
+        pass
+
+    @abstractmethod
+    def create_collection(self, schema, index_params) -> None:
+        pass
+
+    @abstractmethod
+    def insert(self, data) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def upsert(self, data) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def search(
+        self,
+        data,
+        anns_field,
+        limit,
+        filter="",
+        search_params=None,
+        output_fields=None,
+        calculate_recall=False,
+        ground_truth=None,
+    ) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def hybrid_search(self, reqs, ranker, limit, output_fields=None) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def query(self, filter, output_fields=None) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def delete(self, filter) -> dict[str, Any]:
+        pass
+
+
+class MilvusV2Client(BaseClient):
+    """Milvus v2 Python SDK Client Wrapper"""
+
+    def __init__(self, uri, collection_name, token="root:Milvus", db_name="default", timeout=60):
+        self.uri = uri
+        self.collection_name = collection_name
+        self.token = token
+        self.db_name = db_name
+        self.timeout = timeout
+
+        # Initialize MilvusClient v2
+        self.client = MilvusClient(
+            uri=self.uri,
+            token=self.token,
+            db_name=self.db_name,
+            timeout=self.timeout,
+        )
+
+    def close(self):
+        self.client.close()
+
+    def create_collection(self, schema, index_params):
+        self.client.create_collection(
+            collection_name=self.collection_name,
+            schema=schema,
+            index_params=index_params,
+        )
+
+    def insert(self, data):
+        start = time.time()
+        try:
+            result = self.client.insert(collection_name=self.collection_name, data=data)
+            total_time = (time.time() - start) * 1000
+            return {"success": True, "response_time": total_time, "result": result}
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+    def upsert(self, data):
+        start = time.time()
+        try:
+            result = self.client.upsert(collection_name=self.collection_name, data=data)
+            total_time = (time.time() - start) * 1000
+            return {"success": True, "response_time": total_time, "result": result}
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+    def search(
+        self,
+        data,
+        anns_field,
+        limit,
+        filter="",
+        search_params=None,
+        output_fields=None,
+        calculate_recall=False,
+        ground_truth=None,
+    ):
+        if output_fields is None:
+            output_fields = ["id"]
+
+        start = time.time()
+        try:
+            result = self.client.search(
+                collection_name=self.collection_name,
+                data=data,
+                anns_field=anns_field,
+                filter=filter,
+                limit=limit,
+                search_params=search_params,
+                output_fields=output_fields,
+            )
+            total_time = (time.time() - start) * 1000
+            empty = len(result) == 0 or all(len(r) == 0 for r in result)
+
+            # Prepare base result
+            search_result = {
+                "success": not empty,
+                "response_time": total_time,
+                "empty": empty,
+                "result": result,
+            }
+
+            # Calculate recall if requested
+            if calculate_recall and ground_truth is not None and not empty:
+                recall_value = self.get_recall(result, ground_truth, limit)
+                search_result["recall"] = recall_value
+
+            return search_result
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+    def hybrid_search(self, reqs, ranker, limit, output_fields=None):
+        if output_fields is None:
+            output_fields = ["id"]
+
+        start = time.time()
+        try:
+            result = self.client.hybrid_search(
+                collection_name=self.collection_name,
+                reqs=reqs,
+                ranker=ranker,
+                limit=limit,
+                output_fields=output_fields,
+                timeout=self.timeout,
+            )
+            total_time = (time.time() - start) * 1000
+            empty = len(result) == 0 or all(len(r) == 0 for r in result)
+
+            # Prepare base result
+            search_result = {
+                "success": not empty,
+                "response_time": total_time,
+                "empty": empty,
+                "result": result,
+            }
+
+            return search_result
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+    @staticmethod
+    def get_recall(search_results, ground_truth, limit=None):
+        """Calculate recall for V2 client search results."""
+        try:
+            # Extract IDs from V2 search results
+            retrieved_ids = []
+            if isinstance(search_results, list) and len(search_results) > 0:
+                # search_results[0] contains the search results for the first query
+                for hit in search_results[0] if isinstance(search_results[0], list) else search_results:
+                    if isinstance(hit, dict) and "id" in hit:
+                        retrieved_ids.append(hit["id"])
+                    elif hasattr(hit, "get"):
+                        retrieved_ids.append(hit.get("id"))
+
+            # Apply limit if specified
+            if limit is None:
+                limit = len(retrieved_ids)
+
+            if len(ground_truth) < limit:
+                raise ValueError(f"Ground truth length is less than limit: {len(ground_truth)} < {limit}")
+
+            # Calculate recall
+            ground_truth_set = set(ground_truth[:limit])
+            retrieved_set = set(retrieved_ids)
+            intersect = len(ground_truth_set.intersection(retrieved_set))
+            return intersect / len(ground_truth_set)
+
+        except Exception:
+            return 0.0
+
+    def query(self, filter, output_fields=None):
+        if output_fields is None:
+            output_fields = ["id"]
+
+        start = time.time()
+        try:
+            result = self.client.query(
+                collection_name=self.collection_name,
+                filter=filter,
+                output_fields=output_fields,
+            )
+            total_time = (time.time() - start) * 1000
+            empty = len(result) == 0
+            return {
+                "success": not empty,
+                "response_time": total_time,
+                "empty": empty,
+                "result": result,
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+    def delete(self, filter):
+        start = time.time()
+        try:
+            result = self.client.delete(collection_name=self.collection_name, filter=filter)
+            total_time = (time.time() - start) * 1000
+            return {"success": True, "response_time": total_time, "result": result}
+        except Exception as e:
+            return {
+                "success": False,
+                "response_time": (time.time() - start) * 1000,
+                "exception": e,
+            }
+
+
+# ----------------------------------
+# Locust User wrapper
+# ----------------------------------
+
+
+class MilvusUser(User):
+    """Locust User implementation for Milvus operations.
+
+    This class wraps the MilvusV2Client implementation and translates
+    client method results into Locust request events so that performance
+    statistics are collected properly.
+
+    Parameters
+    ----------
+    host : str
+        Milvus server URI, e.g. ``"http://localhost:19530"``.
+    collection_name : str
+        The name of the collection to operate on.
+    **client_kwargs
+        Additional keyword arguments forwarded to the client.
+    """
+
+    abstract = True
+
+    def __init__(
+        self,
+        environment,
+        uri: str = "http://localhost:19530",
+        token: str = "root:Milvus",
+        collection_name: str = "test_collection",
+        db_name: str = "default",
+        timeout: int = 60,
+        schema: CollectionSchema | None = None,
+        index_params: IndexParams | None = None,
+        **client_kwargs,
+    ):
+        super().__init__(environment)
+
+        if uri is None:
+            raise ValueError("'uri' must be provided for MilvusUser")
+        if collection_name is None:
+            raise ValueError("'collection_name' must be provided for MilvusUser")
+
+        self.client_type = "milvus"
+        self.client = MilvusV2Client(
+            uri=uri,
+            token=token,
+            collection_name=collection_name,
+            db_name=db_name,
+            timeout=timeout,
+        )
+        if schema is not None:
+            self.client.create_collection(schema=schema, index_params=index_params)
+
+    @staticmethod
+    def _fire_event(request_type: str, name: str, result: dict[str, Any]):
+        """Emit a Locust request event from a Milvus client result dict."""
+        response_time = int(result.get("response_time", 0))
+        events.request.fire(
+            request_type=f"{request_type}",
+            name=name,
+            response_time=response_time,
+            response_length=0,
+            exception=result.get("exception"),
+        )
+
+    @staticmethod
+    def _fire_recall_event(request_type: str, name: str, result: dict[str, Any]):
+        """Emit a Locust request event for recall metric using recall value instead of response time."""
+        recall_value = result.get("recall", 0.0)
+        # Use recall value as response_time for metric display (scaled by 100 for better visualization) percentage
+        response_time_as_recall = int(recall_value * 100)
+        events.request.fire(
+            request_type=f"{request_type}",
+            name=name,
+            response_time=response_time_as_recall,
+            response_length=result.get("retrieved_count", 0),
+            exception=result.get("exception"),
+        )
+
+    def insert(self, data):
+        result = self.client.insert(data)
+        self._fire_event(self.client_type, "insert", result)
+        return result
+
+    def upsert(self, data):
+        result = self.client.upsert(data)
+        self._fire_event(self.client_type, "upsert", result)
+        return result
+
+    def search(
+        self,
+        data,
+        anns_field,
+        limit,
+        filter="",
+        search_params=None,
+        output_fields=None,
+        calculate_recall=False,
+        ground_truth=None,
+    ):
+        result = self.client.search(
+            data,
+            anns_field,
+            limit,
+            filter=filter,
+            search_params=search_params,
+            output_fields=output_fields,
+            calculate_recall=calculate_recall,
+            ground_truth=ground_truth,
+        )
+        # Fire search event
+        self._fire_event(self.client_type, "search", result)
+
+        # Fire recall event if recall was calculated
+        if calculate_recall and "recall" in result:
+            self._fire_recall_event(self.client_type, "recall", result)
+
+        return result
+
+    def hybrid_search(self, reqs, ranker, limit, output_fields=None):
+        result = self.client.hybrid_search(reqs, ranker, limit, output_fields)
+        self._fire_event(self.client_type, "hybrid_search", result)
+        return result
+
+    def query(self, filter, output_fields=None):
+        result = self.client.query(
+            filter=filter,
+            output_fields=output_fields,
+        )
+        self._fire_event(self.client_type, "query", result)
+        return result
+
+    def delete(self, filter):
+        result = self.client.delete(filter)
+        self._fire_event(self.client_type, "delete", result)
+        return result
+
+    def on_stop(self):
+        self.client.close()

{locust-2.38.1.dev4.dist-info → locust-2.38.2.dev7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: locust
-Version: 2.38.1.dev4
+Version: 2.38.2.dev7
 Summary: Developer-friendly load testing framework
 Project-URL: homepage, https://locust.io/
 Project-URL: repository, https://github.com/locustio/locust
@@ -41,6 +41,8 @@ Requires-Dist: setuptools>=70.0.0
 Requires-Dist: tomli>=1.1.0; python_version < '3.11'
 Requires-Dist: typing-extensions>=4.6.0; python_version < '3.11'
 Requires-Dist: werkzeug>=2.0.0
+Provides-Extra: milvus
+Requires-Dist: pymilvus>=2.5.0; extra == 'milvus'
 Description-Content-Type: text/markdown
 
 # Locust

{locust-2.38.1.dev4.dist-info → locust-2.38.2.dev7.dist-info}/RECORD

@@ -1,6 +1,6 @@
 locust/__init__.py,sha256=HadpgGidiyCDPSKwkxrk1Qw6eB7dTmftNJVftuJzAiw,1876
 locust/__main__.py,sha256=vBQ82334kX06ImDbFlPFgiBRiLIinwNk3z8Khs6hd74,31
-locust/_version.py,sha256=uTueUC_r5h08wNJDaddN1-QEOPEeLyIZpdxIzqXbZ6I,526
+locust/_version.py,sha256=s2j9gE6VWFZ3Bbw6sbjPzaEsKdu7tciZebcEag4lP4c,719
 locust/argument_parser.py,sha256=t6mAoK9u13DxC9UH-alVqS6fFABFTyNWSJG89yQ4QQQ,33056
 locust/clients.py,sha256=o-277lWQdpmPnoRTdf3IQVNPQT8LMFDtPtuxbLHQIIs,19286
 locust/debug.py,sha256=7CCm8bIg44uGH2wqBlo1rXBzV2VzwPicLxLewz8r5CQ,5099
@@ -18,7 +18,8 @@ locust/shape.py,sha256=t-lwBS8LOjWcKXNL7j2U3zroIXJ1b0fazUwpRYQOKXw,1973
 locust/stats.py,sha256=qyoSKT0i7RunLDj5pMGqizK1Sp8bcqUsXwh2m4_DpR8,47203
 locust/web.py,sha256=HLFN9jUtKG3sMIKu_Xw9wtvTAFxXvzDHdtLtfb_JxUQ,31849
 locust/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-locust/contrib/fasthttp.py,sha256=04rnf4kxRxVbzOyqDlb-io43dWIyMPMUwmQ7DOFo2Lo,29656
+locust/contrib/fasthttp.py,sha256=QfVpcdo5oKOO1rKLXeCeTJ4t57XU1lhIIWeUeeSJtIE,29650
+locust/contrib/milvus.py,sha256=YabgLd0lImzWupJFCm0OZAW-Nxeibwn91ldWpZ2irDo,12811
 locust/contrib/mongodb.py,sha256=1seUYgJOaNKwybYOP9PUEVhgl8hGy-G33f8lFj3R8W8,1246
 locust/contrib/oai.py,sha256=Ot3T8lp31ThckGbNps86oVvq6Vn845Eec0mxhDmONDE,2684
 locust/contrib/postgres.py,sha256=OuMWnGYN10K65Tq2axVESEW25Y0g5gJb0rK90jkcCJg,1230
@@ -53,8 +54,8 @@ locust/webui/dist/assets/index-BjqxSg7R.js,sha256=3JyrKWfAg8LlTy2bxAJh73c6njNPhN
 locust/webui/dist/assets/terminal.gif,sha256=iw80LO2u0dnf4wpGfFJZauBeKTcSpw9iUfISXT2nEF4,75302
 locust/webui/dist/assets/testruns-dark.png,sha256=G4p2VZSBuuqF4neqUaPSshIp5OKQJ_Bvb69Luj6XuVs,125231
 locust/webui/dist/assets/testruns-light.png,sha256=JinGDiiBPOkhpfF-XCbmQqhRInqItrjrBTLKt5MlqVI,130301
-locust-2.38.1.dev4.dist-info/METADATA,sha256=LBqunPnbR6lVJHUB6eXzouDXc_60dzInTBZ6hbZpqjo,9403
-locust-2.38.1.dev4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust-2.38.1.dev4.dist-info/entry_points.txt,sha256=RAdt8Ku-56m7bFjmdj-MBhbF6h4NX7tVODR9QNnOg0E,44
-locust-2.38.1.dev4.dist-info/licenses/LICENSE,sha256=5hnz-Vpj0Z3kSCQl0LzV2hT1TLc4LHcbpBp3Cy-EuyM,1110
-locust-2.38.1.dev4.dist-info/RECORD,,
+locust-2.38.2.dev7.dist-info/METADATA,sha256=Y85K_sULv2Nnz46NJgs-rAr3QeH01cT0OOlBsy9LyoA,9476
+locust-2.38.2.dev7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust-2.38.2.dev7.dist-info/entry_points.txt,sha256=RAdt8Ku-56m7bFjmdj-MBhbF6h4NX7tVODR9QNnOg0E,44
+locust-2.38.2.dev7.dist-info/licenses/LICENSE,sha256=5hnz-Vpj0Z3kSCQl0LzV2hT1TLc4LHcbpBp3Cy-EuyM,1110
+locust-2.38.2.dev7.dist-info/RECORD,,