Skip to content

Commit

Permalink
Add db parameter in flush_all (#1701)
Browse files Browse the repository at this point in the history
Signed-off-by: wayblink <[email protected]>
  • Loading branch information
wayblink authored Sep 15, 2023
1 parent 6d9c816 commit c24ed43
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 6 deletions.
4 changes: 2 additions & 2 deletions pymilvus/client/grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions pymilvus/client/prepare.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 2 additions & 0 deletions pymilvus/orm/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down

0 comments on commit c24ed43

Please sign in to comment.