Browse Source

Implement ListAlias interface (#799)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 1 year ago
parent
commit
26e3eafe4e

+ 50 - 12
src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java

@@ -1071,12 +1071,15 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
                 requestParam.getCollectionName(), requestParam.getAlias());
 
         try {
-            CreateAliasRequest createAliasRequest = CreateAliasRequest.newBuilder()
+            CreateAliasRequest.Builder builder = CreateAliasRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
-                    .setAlias(requestParam.getAlias())
-                    .build();
+                    .setAlias(requestParam.getAlias());
+
+            if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
+                builder.setDbName(requestParam.getDatabaseName());
+            }
 
-            Status response = blockingStub().createAlias(createAliasRequest);
+            Status response = blockingStub().createAlias(builder.build());
             handleResponse(title, response);
             return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
         } catch (StatusRuntimeException e) {
@@ -1098,11 +1101,14 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         String title = String.format("DropAliasRequest alias:%s", requestParam.getAlias());
 
         try {
-            DropAliasRequest dropAliasRequest = DropAliasRequest.newBuilder()
-                    .setAlias(requestParam.getAlias())
-                    .build();
+            DropAliasRequest.Builder builder = DropAliasRequest.newBuilder()
+                    .setAlias(requestParam.getAlias());
+
+            if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
+                builder.setDbName(requestParam.getDatabaseName());
+            }
 
-            Status response = blockingStub().dropAlias(dropAliasRequest);
+            Status response = blockingStub().dropAlias(builder.build());
             handleResponse(title, response);
             return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
         } catch (StatusRuntimeException e) {
@@ -1125,12 +1131,15 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
                 requestParam.getCollectionName(), requestParam.getAlias());
 
         try {
-            AlterAliasRequest alterAliasRequest = AlterAliasRequest.newBuilder()
+            AlterAliasRequest.Builder builder = AlterAliasRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
-                    .setAlias(requestParam.getAlias())
-                    .build();
+                    .setAlias(requestParam.getAlias());
+
+            if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
+                builder.setDbName(requestParam.getDatabaseName());
+            }
 
-            Status response = blockingStub().alterAlias(alterAliasRequest);
+            Status response = blockingStub().alterAlias(builder.build());
             handleResponse(title, response);
             return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
         } catch (StatusRuntimeException e) {
@@ -1142,6 +1151,35 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         }
     }
 
+    @Override
+    public R<ListAliasesResponse> listAliases(ListAliasesParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logDebug(requestParam.toString());
+        String title = String.format("ListAliasesRequest collectionName:%s", requestParam.getCollectionName());
+
+        try {
+            ListAliasesRequest.Builder builder = ListAliasesRequest.newBuilder()
+                    .setCollectionName(requestParam.getCollectionName());
+
+            if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
+                builder.setDbName(requestParam.getDatabaseName());
+            }
+
+            ListAliasesResponse response = blockingStub().listAliases(builder.build());
+            handleResponse(title, response.getStatus());
+            return R.success(response);
+        } catch (StatusRuntimeException e) {
+            logError("{} RPC failed! Exception:{}", title, e);
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("{} failed! Exception:{}", title, e);
+            return R.failed(e);
+        }
+    }
+
     @Override
     public R<RpcStatus> createIndex(@NonNull CreateIndexParam requestParam) {
         if (!clientIsReady()) {

+ 8 - 0
src/main/java/io/milvus/client/MilvusClient.java

@@ -308,6 +308,14 @@ public interface MilvusClient {
      */
     R<RpcStatus> alterAlias(AlterAliasParam requestParam);
 
+    /**
+     * List all alias for a collection.
+     *
+     * @param requestParam {@link ListAliasesParam}
+     * @return {status:result code, data:ListAliasesResponse{status, aliases}}
+     */
+    R<ListAliasesResponse> listAliases(ListAliasesParam requestParam);
+
     /**
      * Creates an index on a vector field in the specified collection.
      * Note that index building is an async progress.

+ 8 - 0
src/main/java/io/milvus/client/MilvusMultiServiceClient.java

@@ -308,6 +308,14 @@ public class MilvusMultiServiceClient implements MilvusClient {
         return handleResponse(response);
     }
 
+    @Override
+    public R<ListAliasesResponse> listAliases(ListAliasesParam requestParam) {
+        List<R<ListAliasesResponse>> response = this.clusterFactory.getAvailableServerSettings().stream()
+                .map(serverSetting -> serverSetting.getClient().listAliases(requestParam))
+                .collect(Collectors.toList());
+        return handleResponse(response);
+    }
+
     @Override
     public R<RpcStatus> createIndex(CreateIndexParam requestParam) {
         List<R<RpcStatus>> response = this.clusterFactory.getAvailableServerSettings().parallelStream()

+ 15 - 0
src/main/java/io/milvus/param/alias/AlterAliasParam.java

@@ -12,10 +12,12 @@ import lombok.NonNull;
 @Getter
 public class AlterAliasParam {
     private final String collectionName;
+    private final String databaseName;
     private final String alias;
 
     private AlterAliasParam(@NonNull AlterAliasParam.Builder builder) {
         this.collectionName = builder.collectionName;
+        this.databaseName = builder.databaseName;
         this.alias = builder.alias;
     }
 
@@ -28,6 +30,7 @@ public class AlterAliasParam {
      */
     public static final class Builder {
         private String collectionName;
+        private String databaseName;
         private String alias;
 
         private Builder() {
@@ -44,6 +47,17 @@ public class AlterAliasParam {
             return this;
         }
 
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(String databaseName) {
+            this.databaseName = databaseName;
+            return this;
+        }
+
         /**
          * Sets the collection alias. Collection alias cannot be empty or null.
          *
@@ -77,6 +91,7 @@ public class AlterAliasParam {
     public String toString() {
         return "AlterAliasParam{" +
                 "collectionName='" + collectionName + '\'' +
+                "databaseName='" + databaseName + '\'' +
                 ", alias='" + alias + '\'' +
                 '}';
     }

+ 15 - 0
src/main/java/io/milvus/param/alias/CreateAliasParam.java

@@ -12,10 +12,12 @@ import lombok.NonNull;
 @Getter
 public class CreateAliasParam {
     private final String collectionName;
+    private final String databaseName;
     private final String alias;
 
     private CreateAliasParam(@NonNull CreateAliasParam.Builder builder) {
         this.collectionName = builder.collectionName;
+        this.databaseName = builder.databaseName;
         this.alias = builder.alias;
     }
 
@@ -28,6 +30,7 @@ public class CreateAliasParam {
      */
     public static final class Builder {
         private String collectionName;
+        private String databaseName;
         private String alias;
 
         private Builder() {
@@ -44,6 +47,17 @@ public class CreateAliasParam {
             return this;
         }
 
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(String databaseName) {
+            this.databaseName = databaseName;
+            return this;
+        }
+
         /**
          * Sets the collection alias. Collection alias cannot be empty or null.
          *
@@ -77,6 +91,7 @@ public class CreateAliasParam {
     public String toString() {
         return "CreateAliasParam{" +
                 "collectionName='" + collectionName + '\'' +
+                "databaseName='" + databaseName + '\'' +
                 ", alias='" + alias + '\'' +
                 '}';
     }

+ 15 - 0
src/main/java/io/milvus/param/alias/DropAliasParam.java

@@ -12,9 +12,11 @@ import lombok.NonNull;
 @Getter
 public class DropAliasParam {
     private final String alias;
+    private final String databaseName;
 
     private DropAliasParam(@NonNull Builder builder) {
         this.alias = builder.alias;
+        this.databaseName = builder.databaseName;
     }
 
     public static Builder newBuilder() {
@@ -26,6 +28,7 @@ public class DropAliasParam {
      */
     public static final class Builder {
         private String alias;
+        private String databaseName;
 
         private Builder() {
         }
@@ -41,6 +44,17 @@ public class DropAliasParam {
             return this;
         }
 
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(String databaseName) {
+            this.databaseName = databaseName;
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link DropAliasParam} instance.
          *
@@ -62,6 +76,7 @@ public class DropAliasParam {
     public String toString() {
         return "DropAliasParam{" +
                 ", alias='" + alias + '\'' +
+                "databaseName='" + databaseName + '\'' +
                 '}';
     }
 }

+ 79 - 0
src/main/java/io/milvus/param/alias/ListAliasesParam.java

@@ -0,0 +1,79 @@
+package io.milvus.param.alias;
+
+import io.milvus.exception.ParamException;
+import io.milvus.param.ParamUtils;
+
+import lombok.Getter;
+import lombok.NonNull;
+
+@Getter
+public class ListAliasesParam {
+    private final String collectionName;
+    private final String databaseName;
+
+    private ListAliasesParam(@NonNull Builder builder) {
+        this.collectionName = builder.collectionName;
+        this.databaseName = builder.databaseName;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Builder for {@link ListAliasesParam} class.
+     */
+    public static final class Builder {
+        private String collectionName;
+        private String databaseName;
+
+        private Builder() {
+        }
+
+        /**
+         * Sets the collection name. Collection name cannot be empty or null.
+         *
+         * @param collectionName collection name
+         * @return <code>Builder</code>
+         */
+        public Builder withCollectionName(@NonNull String collectionName) {
+            this.collectionName = collectionName;
+            return this;
+        }
+
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(String databaseName) {
+            this.databaseName = databaseName;
+            return this;
+        }
+
+        /**
+         * Verifies parameters and creates a new {@link ListAliasesParam} instance.
+         *
+         * @return {@link ListAliasesParam}
+         */
+        public ListAliasesParam build() throws ParamException {
+            ParamUtils.CheckNullEmptyString(collectionName, "Collection name");
+
+            return new ListAliasesParam(this);
+        }
+    }
+
+    /**
+     * Constructs a <code>String</code> by {@link ListAliasesParam} instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "ListAliasesParam{" +
+                "collectionName='" + collectionName + '\'' +
+                "databaseName='" + databaseName + '\'' +
+                '}';
+    }
+}

+ 17 - 4
src/main/java/io/milvus/response/DescCollResponseWrapper.java

@@ -1,16 +1,15 @@
 package io.milvus.response;
 
 import io.milvus.exception.ParamException;
-import io.milvus.grpc.CollectionSchema;
-import io.milvus.grpc.DataType;
-import io.milvus.grpc.DescribeCollectionResponse;
-import io.milvus.grpc.FieldSchema;
+import io.milvus.grpc.*;
 import io.milvus.param.ParamUtils;
 import io.milvus.param.collection.FieldType;
 import lombok.NonNull;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Util class to wrap response of <code>describeCollection</code> interface.
@@ -185,6 +184,19 @@ public class DescCollResponseWrapper {
         throw new ParamException("No vector key found.");
     }
 
+    /**
+     * Get properties of the collection.
+     *
+     * @return List of String, aliases of the collection
+     */
+    public Map<String, String> getProperties() {
+        Map<String, String> pairs = new HashMap<>();
+        List<KeyValuePair> props = response.getPropertiesList();
+        props.forEach((prop) -> pairs.put(prop.getKey(), prop.getValue()));
+
+        return pairs;
+    }
+
     /**
      * Construct a <code>String</code> by {@link DescCollResponseWrapper} instance.
      *
@@ -201,6 +213,7 @@ public class DescCollResponseWrapper {
                 ", aliases:" + getAliases().toString() +
                 ", fields:" + getFields().toString() +
                 ", isDynamicFieldEnabled:" + isDynamicFieldEnabled() +
+                ", properties:" + getProperties() +
                 '}';
     }
 }

+ 87 - 0
src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -24,6 +24,10 @@ import com.google.common.util.concurrent.ListenableFuture;
 import io.milvus.common.clientenum.ConsistencyLevelEnum;
 import io.milvus.grpc.*;
 import io.milvus.param.*;
+import io.milvus.param.alias.AlterAliasParam;
+import io.milvus.param.alias.CreateAliasParam;
+import io.milvus.param.alias.DropAliasParam;
+import io.milvus.param.alias.ListAliasesParam;
 import io.milvus.param.collection.*;
 import io.milvus.param.dml.InsertParam;
 import io.milvus.param.dml.QueryParam;
@@ -1846,6 +1850,89 @@ class MilvusClientDockerTest {
         Assertions.assertEquals(R.Status.Success.getCode(), dropR.getStatus().intValue());
     }
 
+    @Test
+    void testAlias() {
+        // collection schema
+        List<FieldType> fieldsSchema = new ArrayList<>();
+        fieldsSchema.add(FieldType.newBuilder()
+                .withPrimaryKey(true)
+                .withDataType(DataType.Int64)
+                .withName("id")
+                .build());
+
+        fieldsSchema.add(FieldType.newBuilder()
+                .withDataType(DataType.FloatVector)
+                .withName("vector")
+                .withDimension(dimension)
+                .build());
+
+        // create collection A
+        R<RpcStatus> response = client.createCollection(CreateCollectionParam.newBuilder()
+                .withCollectionName("coll_A")
+                .withFieldTypes(fieldsSchema)
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), response.getStatus().intValue());
+
+        // create collection B
+        response = client.createCollection(CreateCollectionParam.newBuilder()
+                .withCollectionName("coll_B")
+                .withFieldTypes(fieldsSchema)
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), response.getStatus().intValue());
+
+        // create alias
+        response = client.createAlias(CreateAliasParam.newBuilder()
+                .withCollectionName("coll_A")
+                .withAlias("alias_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), response.getStatus().intValue());
+
+        R<Boolean> has = client.hasCollection(HasCollectionParam.newBuilder()
+                .withCollectionName("alias_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), has.getStatus().intValue());
+        Assertions.assertEquals(has.getData(), true);
+
+        R<ListAliasesResponse> listResp = client.listAliases(ListAliasesParam.newBuilder()
+                .withCollectionName("coll_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), listResp.getStatus().intValue());
+        Assertions.assertEquals(listResp.getData().getAliases(0), "alias_A");
+        Assertions.assertEquals(listResp.getData().getCollectionName(), "coll_A");
+
+        // alter alias
+        response = client.alterAlias(AlterAliasParam.newBuilder()
+                .withAlias("alias_A")
+                .withCollectionName("coll_B")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), listResp.getStatus().intValue());
+
+        has = client.hasCollection(HasCollectionParam.newBuilder()
+                .withCollectionName("alias_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), has.getStatus().intValue());
+        Assertions.assertEquals(has.getData(), true);
+
+        listResp = client.listAliases(ListAliasesParam.newBuilder()
+                .withCollectionName("coll_B")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), listResp.getStatus().intValue());
+        Assertions.assertEquals(listResp.getData().getAliases(0), "alias_A");
+        Assertions.assertEquals(listResp.getData().getCollectionName(), "coll_B");
+
+        // drop alias
+        response = client.dropAlias(DropAliasParam.newBuilder()
+                .withAlias("alias_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), response.getStatus().intValue());
+
+        has = client.hasCollection(HasCollectionParam.newBuilder()
+                .withCollectionName("alias_A")
+                .build());
+        Assertions.assertEquals(R.Status.Success.getCode(), has.getStatus().intValue());
+        Assertions.assertEquals(has.getData(), false);
+    }
+
     @Test
     void testHighLevelGet() {
         // collection schema

+ 40 - 2
src/test/java/io/milvus/client/MilvusServiceClientTest.java

@@ -29,6 +29,7 @@ import io.milvus.param.*;
 import io.milvus.param.alias.AlterAliasParam;
 import io.milvus.param.alias.CreateAliasParam;
 import io.milvus.param.alias.DropAliasParam;
+import io.milvus.param.alias.ListAliasesParam;
 import io.milvus.param.collection.*;
 import io.milvus.param.control.*;
 import io.milvus.param.credential.*;
@@ -1197,6 +1198,12 @@ class MilvusServiceClientTest {
                 .withAlias("")
                 .build()
         );
+
+        CreateAliasParam param = CreateAliasParam.newBuilder()
+                .withCollectionName("collection1")
+                .withAlias("alias1")
+                .withDatabaseName("db1")
+                .build();
     }
 
     @Test
@@ -1222,6 +1229,7 @@ class MilvusServiceClientTest {
     void dropAlias() {
         DropAliasParam param = DropAliasParam.newBuilder()
                 .withAlias("alias1")
+                .withDatabaseName("db1")
                 .build();
 
         testFuncByName("dropAlias", param);
@@ -1230,17 +1238,23 @@ class MilvusServiceClientTest {
     @Test
     void alterAliasParam() {
         // test throw exception with illegal input
-        assertThrows(ParamException.class, () -> CreateAliasParam.newBuilder()
+        assertThrows(ParamException.class, () -> AlterAliasParam.newBuilder()
                 .withCollectionName("")
                 .withAlias("alias1")
                 .build()
         );
 
-        assertThrows(ParamException.class, () -> CreateAliasParam.newBuilder()
+        assertThrows(ParamException.class, () -> AlterAliasParam.newBuilder()
                 .withCollectionName("collection1")
                 .withAlias("")
                 .build()
         );
+
+        AlterAliasParam param = AlterAliasParam.newBuilder()
+                .withCollectionName("collection1")
+                .withAlias("alias1")
+                .withDatabaseName("db1")
+                .build();
     }
 
     @Test
@@ -1253,6 +1267,30 @@ class MilvusServiceClientTest {
         testFuncByName("alterAlias", param);
     }
 
+    @Test
+    void listAliasesParam() {
+        // test throw exception with illegal input
+        assertThrows(ParamException.class, () -> ListAliasesParam.newBuilder()
+                .withCollectionName("")
+                .withDatabaseName("")
+                .build()
+        );
+
+        ListAliasesParam param = ListAliasesParam.newBuilder()
+                .withCollectionName("collection1")
+                .withDatabaseName("")
+                .build();
+    }
+
+    @Test
+    void listAliases() {
+        ListAliasesParam param = ListAliasesParam.newBuilder()
+                .withCollectionName("collection1")
+                .build();
+
+        testFuncByName("listAliases", param);
+    }
+
     @Test
     void createIndexParam() {
         // test throw exception with illegal input

+ 14 - 0
src/test/java/io/milvus/server/MockMilvusServerImpl.java

@@ -45,6 +45,7 @@ public class MockMilvusServerImpl extends MilvusServiceGrpc.MilvusServiceImplBas
     private io.milvus.grpc.Status respCreateAlias;
     private io.milvus.grpc.Status respDropAlias;
     private io.milvus.grpc.Status respAlterAlias;
+    private io.milvus.grpc.ListAliasesResponse respListAliases;
     private io.milvus.grpc.Status respCreateIndex;
     private io.milvus.grpc.DescribeIndexResponse respDescribeIndex;
     private io.milvus.grpc.GetIndexStateResponse respGetIndexState;
@@ -324,6 +325,19 @@ public class MockMilvusServerImpl extends MilvusServiceGrpc.MilvusServiceImplBas
         respAlterAlias = resp;
     }
 
+    @Override
+    public void listAliases(io.milvus.grpc.ListAliasesRequest request,
+                           io.grpc.stub.StreamObserver<io.milvus.grpc.ListAliasesResponse> responseObserver) {
+        logger.info("MockServer receive listAliases() call");
+
+        responseObserver.onNext(respListAliases);
+        responseObserver.onCompleted();
+    }
+
+    public void setListAliasesResponse(io.milvus.grpc.ListAliasesResponse resp) {
+        respListAliases = resp;
+    }
+
     @Override
     public void createIndex(io.milvus.grpc.CreateIndexRequest request,
                             io.grpc.stub.StreamObserver<io.milvus.grpc.Status> responseObserver) {