Add new CLI commands for collection management and data op

milvus_cli/Collection.py

@@ -206,3 +206,81 @@ def list_fields_info(self, collectionName):
             }
             for i in result
         ]
+
+    def flush(self, collectionName, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            target.flush(timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Flush collection error!{str(e)}")
+        else:
+            return f"Flush collection {collectionName} successfully!"
+
+    def compact(self, collectionName, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            compaction_id = target.compact(timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Compact collection error!{str(e)}")
+        else:
+            return {"message": f"Compact collection {collectionName} successfully!", "compaction_id": compaction_id}
+
+    def get_compaction_state(self, collectionName, compactionId, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            state = target.get_compaction_state(compactionId, timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Get compaction state error!{str(e)}")
+        else:
+            return state
+
+    def wait_for_compaction_completed(self, collectionName, compactionId, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            target.wait_for_compaction_completed(compactionId, timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Wait for compaction completed error!{str(e)}")
+        else:
+            return f"Compaction {compactionId} completed!"
+
+    def get_compaction_plans(self, collectionName, compactionId, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            plans = target.get_compaction_plans(compactionId, timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Get compaction plans error!{str(e)}")
+        else:
+            return plans
+
+    def get_replicas(self, collectionName, timeout=None):
+        try:
+            target = getTargetCollection(collectionName)
+            replicas = target.get_replicas(timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Get replicas error!{str(e)}")
+        else:
+            return replicas
+
+    def truncate(self, collectionName, timeout=None):
+        try:
+            utility.truncate_collection(collectionName, timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Truncate collection error!{str(e)}")
+        else:
+            return f"Truncate collection {collectionName} successfully!"
+
+    def load_state(self, collectionName, partitionName=None):
+        try:
+            state = utility.load_state(collectionName, partitionName)
+        except Exception as e:
+            raise Exception(f"Get load state error!{str(e)}")
+        else:
+            return state
+
+    def wait_for_loading_complete(self, collectionName, partitionNames=None, timeout=None):
+        try:
+            utility.wait_for_loading_complete(collectionName, partitionNames, timeout=timeout)
+        except Exception as e:
+            raise Exception(f"Wait for loading complete error!{str(e)}")
+        else:
+            return f"Loading {collectionName} completed!"

milvus_cli/Data.py

@@ -33,6 +33,13 @@ def delete_entities(
         result = collection.delete(expr, partition_name=partition_name)
         return result
 
+    def upsert(self, collectionName, data, partitionName=None, timeout=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        result = collection.upsert(data, partition_name=partitionName, timeout=timeout)
+        return result
+
     def search(self, collectionName, searchParameters, prettierFormat=True):
         collection = getTargetCollection(
             collectionName,
@@ -47,3 +54,80 @@ def search(self, collectionName, searchParameters, prettierFormat=True):
             tablefmt="grid",
         )
         return results
+
+    def hybrid_search(self, collectionName, requests, rerank, limit, output_fields=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        res = collection.hybrid_search(requests, rerank, limit, output_fields=output_fields)
+        # Format results
+        if res and len(res) > 0:
+            hits = res[0]
+            results = tabulate(
+                [[hit.id, hit.distance, str(hit.entity)] for hit in hits],
+                headers=["ID", "Distance", "Entity"],
+                tablefmt="grid",
+            )
+            return results
+        return "No results found."
+
+    def query_iterator(self, collectionName, batch_size, expr=None, output_fields=None, partition_names=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        return collection.query_iterator(
+            batch_size=batch_size,
+            expr=expr,
+            output_fields=output_fields,
+            partition_names=partition_names,
+        )
+
+    def search_iterator(self, collectionName, batch_size, data, ann_field, param, limit, expr=None, output_fields=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        return collection.search_iterator(
+            data=data,
+            ann_field=ann_field,
+            param=param,
+            batch_size=batch_size,
+            limit=limit,
+            expr=expr,
+            output_fields=output_fields,
+        )
+
+    def get_by_ids(self, collectionName, ids, output_fields=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        expr = f"{collection.primary_field.name} in {ids}"
+        res = collection.query(expr=expr, output_fields=output_fields)
+        return res
+
+    def delete_by_ids(self, collectionName, ids, partition_name=None):
+        collection = getTargetCollection(
+            collectionName,
+        )
+        expr = f"{collection.primary_field.name} in {ids}"
+        result = collection.delete(expr, partition_name=partition_name)
+        return result
+
+    def bulk_insert(self, collectionName, files, partition_name=None):
+        from pymilvus import utility
+
+        task_id = utility.do_bulk_insert(
+            collection_name=collectionName, partition_name=partition_name, files=files
+        )
+        return task_id
+
+    def get_bulk_insert_state(self, task_id):
+        from pymilvus import utility
+
+        state = utility.get_bulk_insert_state(task_id)
+        return state
+
+    def list_bulk_insert_tasks(self, limit=None, collectionName=None):
+        from pymilvus import utility
+
+        tasks = utility.list_bulk_insert_tasks(limit=limit, collection_name=collectionName)
+        return tasks

milvus_cli/Validation.py

@@ -110,6 +110,7 @@ def validateSearchParams(
     roundDecimal,
     hasIndex=True,
     guarantee_timestamp=None,
+    partitionNames=None,
 ):
     import json
 
@@ -191,6 +192,11 @@ def validateSearchParams(
     else:
         nameList = outputFields.replace(" ", "").split(",")
         result["output_fields"] = nameList
+    if not partitionNames:
+        result["partition_names"] = None
+    else:
+        nameList = partitionNames.replace(" ", "").split(",")
+        result["partition_names"] = nameList
     return result