From c24ed43d7bb201a53acdccf83b9a808303ebcba7 Mon Sep 17 00:00:00 2001
From: wayblink <anyang.wang@zilliz.com>
Date: Fri, 15 Sep 2023 17:29:15 +0800
Subject: [PATCH] Add db parameter in flush_all (#1701)

Signed-off-by: wayblink <anyang.wang@zilliz.com>
---
 pymilvus/client/grpc_handler.py | 4 ++--
 pymilvus/client/prepare.py      | 8 ++++----
 pymilvus/orm/utility.py         | 2 ++
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
index a1a69befd..22a8a5bc9 100644
--- a/pymilvus/client/grpc_handler.py
+++ b/pymilvus/client/grpc_handler.py
@@ -1804,7 +1804,7 @@ def transfer_replica(
 
     @retry_on_rpc_failure()
     def get_flush_all_state(self, flush_all_ts: int, timeout: Optional[float] = None, **kwargs):
-        req = Prepare.get_flush_all_state_request(flush_all_ts)
+        req = Prepare.get_flush_all_state_request(flush_all_ts, kwargs.get("db", ""))
         response = self._stub.GetFlushAllState(req, timeout=timeout)
         status = response.status
         if status.error_code == 0:
@@ -1827,7 +1827,7 @@ def _wait_for_flush_all(self, flush_all_ts: int, timeout: Optional[float] = None
 
     @retry_on_rpc_failure()
     def flush_all(self, timeout: Optional[float] = None, **kwargs):
-        request = Prepare.flush_all_request()
+        request = Prepare.flush_all_request(kwargs.get("db", ""))
         future = self._stub.FlushAll.future(request, timeout=timeout)
         response = future.result()
         if response.status.error_code != 0:
diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
index f6e9f2bad..5a05c2bc8 100644
--- a/pymilvus/client/prepare.py
+++ b/pymilvus/client/prepare.py
@@ -1092,12 +1092,12 @@ def transfer_replica(cls, source: str, target: str, collection_name: str, num_re
         )
 
     @classmethod
-    def flush_all_request(cls):
-        return milvus_types.FlushAllRequest()
+    def flush_all_request(cls, db_name: str):
+        return milvus_types.FlushAllRequest(db_name=db_name)
 
     @classmethod
-    def get_flush_all_state_request(cls, flush_all_ts: int):
-        return milvus_types.GetFlushAllStateRequest(flush_all_ts=flush_all_ts)
+    def get_flush_all_state_request(cls, flush_all_ts: int, db_name: str):
+        return milvus_types.GetFlushAllStateRequest(flush_all_ts=flush_all_ts, db_name=db_name)
 
     @classmethod
     def register_request(cls, user: str, host: str, **kwargs):
diff --git a/pymilvus/orm/utility.py b/pymilvus/orm/utility.py
index affcff3b4..675ab0e74 100644
--- a/pymilvus/orm/utility.py
+++ b/pymilvus/orm/utility.py
@@ -1193,6 +1193,8 @@ def flush_all(using: str = "default", timeout: Optional[float] = None, **kwargs)
             If timeout is not set, the client keeps waiting until the server responds or
             an error occurs.
         **kwargs (``dict``, optional):
+            * *db*(``string``)
+                database to flush.
             * *_async*(``bool``)
                 Indicate if invoke asynchronously. Default `False`.