diff --git a/client/client_grpc_admin.go b/client/admin.go similarity index 94% rename from client/client_grpc_admin.go rename to client/admin.go index 49527dec..8f0dc6c7 100644 --- a/client/client_grpc_admin.go +++ b/client/admin.go @@ -17,6 +17,7 @@ import ( server "github.com/milvus-io/milvus-proto/go-api/milvuspb" ) +// GetVersion returns milvus server version information. func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) { if c.Service == nil { return "", ErrClientNotReady diff --git a/client/client_grpc_alias.go b/client/alias.go similarity index 100% rename from client/client_grpc_alias.go rename to client/alias.go diff --git a/client/client_grpc_alias_test.go b/client/alias_test.go similarity index 100% rename from client/client_grpc_alias_test.go rename to client/alias_test.go diff --git a/client/client_grpc_authentication.go b/client/authentication.go similarity index 100% rename from client/client_grpc_authentication.go rename to client/authentication.go diff --git a/client/client_grpc_authentication_test.go b/client/authentication_test.go similarity index 100% rename from client/client_grpc_authentication_test.go rename to client/authentication_test.go diff --git a/client/client_test.go b/client/client_test.go index c8c70468..89ad7bbf 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -47,50 +47,15 @@ const ( ) func defaultSchema() *entity.Schema { - return &entity.Schema{ - CollectionName: testCollectionName, - AutoID: false, - Fields: []*entity.Field{ - { - Name: testPrimaryField, - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - AutoID: true, - }, - { - Name: testVectorField, - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{ - entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim), - }, - }, - }, - } + return entity.NewSchema().WithName(testCollectionName).WithAutoID(false). + WithField(entity.NewField().WithName(testPrimaryField).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)). + WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim)) } func varCharSchema() *entity.Schema { - return &entity.Schema{ - CollectionName: testCollectionName, - AutoID: false, - Fields: []*entity.Field{ - { - Name: "varchar", - DataType: entity.FieldTypeVarChar, - PrimaryKey: true, - AutoID: false, - TypeParams: map[string]string{ - entity.TypeParamMaxLength: fmt.Sprintf("%d", 100), - }, - }, - { - Name: testVectorField, - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{ - entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim), - }, - }, - }, - } + return entity.NewSchema().WithName(testCollectionName).WithAutoID(false). + WithField(entity.NewField().WithName("varchar").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(false).WithMaxLength(100)). + WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim)) } var _ entity.Row = &defaultRow{} diff --git a/client/collection.go b/client/collection.go index 4d2b5a49..1cd9e4cc 100644 --- a/client/collection.go +++ b/client/collection.go @@ -20,43 +20,11 @@ import ( "github.com/golang/protobuf/proto" "github.com/milvus-io/milvus-sdk-go/v2/entity" - "google.golang.org/grpc" common "github.com/milvus-io/milvus-proto/go-api/commonpb" server "github.com/milvus-io/milvus-proto/go-api/milvuspb" ) -// GrpcClient, uses default grpc Service definition to connect with Milvus2.0 -type GrpcClient struct { - Conn *grpc.ClientConn // grpc connection instance - Service server.MilvusServiceClient // Service client stub -} - -// connect connect to Service -func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error { - if addr == "" { - return fmt.Errorf("address is empty") - } - conn, err := grpc.DialContext(ctx, addr, opts...) - if err != nil { - return err - } - - c.Conn = conn - c.Service = server.NewMilvusServiceClient(c.Conn) - return nil -} - -// Close close the connection -func (c *GrpcClient) Close() error { - if c.Conn != nil { - err := c.Conn.Close() - c.Conn = nil - return err - } - return nil -} - // handles response status // if status is nil returns ErrStatusNil // if status.ErrorCode is common.ErrorCode_Success, returns nil @@ -222,7 +190,7 @@ func (c *GrpcClient) DescribeCollection(ctx context.Context, collName string) (* collection := &entity.Collection{ ID: resp.GetCollectionID(), Name: collName, - Schema: (&entity.Schema{}).ReadProto(resp.GetSchema()), + Schema: entity.NewSchema().ReadProto(resp.GetSchema()), PhysicalChannels: resp.GetPhysicalChannelNames(), VirtualChannels: resp.GetVirtualChannelNames(), ConsistencyLevel: entity.ConsistencyLevel(resp.ConsistencyLevel), diff --git a/client/collection_test.go b/client/collection_test.go index 60e23cad..29104465 100644 --- a/client/collection_test.go +++ b/client/collection_test.go @@ -166,97 +166,34 @@ func (s *CollectionSuite) TestCreateCollection() { } cases := []testCase{ { - name: "empty_fields", - schema: &entity.Schema{ - CollectionName: testCollectionName, - Fields: []*entity.Field{}, - }, + name: "empty_fields", + schema: entity.NewSchema().WithName(testCollectionName), }, { name: "empty_collection_name", - schema: &entity.Schema{ - CollectionName: "", - Fields: []*entity.Field{ - { - Name: "int64", - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - }, - { - Name: "vector", - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{entity.TypeParamDim: "128"}, - }, - }, - }, + schema: entity.NewSchema(). + WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)). + WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)), }, - { name: "multiple primary key", - schema: &entity.Schema{ - - CollectionName: testCollectionName, - Fields: []*entity.Field{ - { - Name: "int64", - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - }, - { - Name: "int64_2", - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - }, - { - Name: "vector", - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{entity.TypeParamDim: "128"}, - }, - }, - }, + schema: entity.NewSchema().WithName(testCollectionName). + WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)). + WithField(entity.NewField().WithName("int64_2").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)). + WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)), }, { name: "multiple auto id", - schema: &entity.Schema{ - CollectionName: testCollectionName, - Fields: []*entity.Field{ - { - Name: "int64", - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - AutoID: true, - }, - { - Name: "int64_2", - DataType: entity.FieldTypeInt64, - PrimaryKey: false, - AutoID: true, - }, - { - Name: "vector", - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{entity.TypeParamDim: "128"}, - }, - }, - }, + schema: entity.NewSchema().WithName(testCollectionName). + WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)). + WithField(entity.NewField().WithName("int64_2").WithDataType(entity.FieldTypeInt64).WithIsAutoID(true)). + WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)), }, { name: "bad_pk_type", - schema: &entity.Schema{ - CollectionName: testCollectionName, - Fields: []*entity.Field{ - { - Name: "float_pk", - DataType: entity.FieldTypeFloat, - PrimaryKey: true, - }, - { - Name: "vector", - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{entity.TypeParamDim: "128"}, - }, - }, - }, + schema: entity.NewSchema(). + WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeDouble).WithIsPrimaryKey(true)). + WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)), }, } diff --git a/client/connect.go b/client/connect.go new file mode 100644 index 00000000..60be44f2 --- /dev/null +++ b/client/connect.go @@ -0,0 +1,51 @@ +// Copyright (C) 2019-2021 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +package client + +import ( + "context" + "fmt" + + server "github.com/milvus-io/milvus-proto/go-api/milvuspb" + "google.golang.org/grpc" +) + +// GrpcClient, uses default grpc Service definition to connect with Milvus2.0 +type GrpcClient struct { + Conn *grpc.ClientConn // grpc connection instance + Service server.MilvusServiceClient // Service client stub +} + +// connect connect to Service +func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error { + if addr == "" { + return fmt.Errorf("address is empty") + } + conn, err := grpc.DialContext(ctx, addr, opts...) + if err != nil { + return err + } + + c.Conn = conn + c.Service = server.NewMilvusServiceClient(c.Conn) + return nil +} + +// Close close the connection +func (c *GrpcClient) Close() error { + if c.Conn != nil { + err := c.Conn.Close() + c.Conn = nil + return err + } + return nil +} diff --git a/client/data_test.go b/client/data_test.go index efcaf7d6..6ce47c7e 100644 --- a/client/data_test.go +++ b/client/data_test.go @@ -13,7 +13,6 @@ package client import ( "context" - "fmt" "math/rand" "testing" "time" @@ -954,56 +953,17 @@ func TestIsCollectionPrimaryKey(t *testing.T) { func TestEstRowSize(t *testing.T) { // a schema contains all supported vector - sch := &entity.Schema{ - CollectionName: testCollectionName, - AutoID: false, - Fields: []*entity.Field{ - { - Name: testPrimaryField, - DataType: entity.FieldTypeInt64, - PrimaryKey: true, - AutoID: true, - }, - { - Name: "attr1", - DataType: entity.FieldTypeInt8, - }, - { - Name: "attr2", - DataType: entity.FieldTypeInt16, - }, - { - Name: "attr3", - DataType: entity.FieldTypeInt32, - }, - { - Name: "attr4", - DataType: entity.FieldTypeFloat, - }, - { - Name: "attr5", - DataType: entity.FieldTypeDouble, - }, - { - Name: "attr6", - DataType: entity.FieldTypeBool, - }, - { - Name: testVectorField, - DataType: entity.FieldTypeFloatVector, - TypeParams: map[string]string{ - entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim), - }, - }, - { - Name: "binary_vector", - DataType: entity.FieldTypeBinaryVector, - TypeParams: map[string]string{ - entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim), - }, - }, - }, - } + sch := entity.NewSchema().WithName(testCollectionName).WithAutoID(false). + WithField(entity.NewField().WithName(testPrimaryField).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)). + WithField(entity.NewField().WithName("attr1").WithDataType(entity.FieldTypeInt8)). + WithField(entity.NewField().WithName("attr2").WithDataType(entity.FieldTypeInt16)). + WithField(entity.NewField().WithName("attr3").WithDataType(entity.FieldTypeInt32)). + WithField(entity.NewField().WithName("attr4").WithDataType(entity.FieldTypeFloat)). + WithField(entity.NewField().WithName("attr5").WithDataType(entity.FieldTypeDouble)). + WithField(entity.NewField().WithName("attr6").WithDataType(entity.FieldTypeBool)). + WithField(entity.NewField().WithName("attr6").WithDataType(entity.FieldTypeBool)). + WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim)). + WithField(entity.NewField().WithName("binary_vector").WithDataType(entity.FieldTypeBinaryVector).WithDim(testVectorDim)) // one row columnID := entity.NewColumnInt64(testPrimaryField, []int64{0}) diff --git a/client/client_grpc_index.go b/client/index.go similarity index 100% rename from client/client_grpc_index.go rename to client/index.go diff --git a/client/client_grpc_index_test.go b/client/index_test.go similarity index 100% rename from client/client_grpc_index_test.go rename to client/index_test.go diff --git a/client/client_grpc_control.go b/client/maintainance.go similarity index 100% rename from client/client_grpc_control.go rename to client/maintainance.go diff --git a/client/client_grpc_control_test.go b/client/maintainance_test.go similarity index 100% rename from client/client_grpc_control_test.go rename to client/maintainance_test.go diff --git a/client/client_grpc_options.go b/client/options.go similarity index 100% rename from client/client_grpc_options.go rename to client/options.go diff --git a/client/client_grpc_options_test.go b/client/options_test.go similarity index 100% rename from client/client_grpc_options_test.go rename to client/options_test.go diff --git a/client/client_grpc_partition.go b/client/partition.go similarity index 100% rename from client/client_grpc_partition.go rename to client/partition.go diff --git a/client/client_grpc_partition_test.go b/client/partition_test.go similarity index 100% rename from client/client_grpc_partition_test.go rename to client/partition_test.go diff --git a/client/client_grpc_rbac.go b/client/rbac.go similarity index 100% rename from client/client_grpc_rbac.go rename to client/rbac.go diff --git a/client/client_grpc_rbac_test.go b/client/rbac_test.go similarity index 100% rename from client/client_grpc_rbac_test.go rename to client/rbac_test.go diff --git a/client/client_grpc_rg.go b/client/resource_group.go similarity index 100% rename from client/client_grpc_rg.go rename to client/resource_group.go diff --git a/client/client_grpc_rg_test.go b/client/resource_group_test.go similarity index 100% rename from client/client_grpc_rg_test.go rename to client/resource_group_test.go diff --git a/entity/schema.go b/entity/schema.go index ca564577..8b7810b1 100644 --- a/entity/schema.go +++ b/entity/schema.go @@ -220,6 +220,14 @@ func (f *Field) WithDim(dim int64) *Field { return f } +func (f *Field) WithMaxLength(maxLen int64) *Field { + if f.TypeParams == nil { + f.TypeParams = make(map[string]string) + } + f.TypeParams[TypeParamMaxLength] = strconv.FormatInt(maxLen, 10) + return f +} + // ReadProto parses FieldSchema func (f *Field) ReadProto(p *schema.FieldSchema) *Field { f.ID = p.GetFieldID() diff --git a/entity/schema_test.go b/entity/schema_test.go index c4efd9c3..b6f55fc5 100644 --- a/entity/schema_test.go +++ b/entity/schema_test.go @@ -65,12 +65,14 @@ func (s *SchemaSuite) TestBasic() { { "test_collection", NewSchema().WithName("test_collection_1").WithDescription("test_collection_1 desc").WithAutoID(false). - WithField(NewField().WithName("ID").WithDataType(FieldTypeInt64).WithIsPrimaryKey(true)), + WithField(NewField().WithName("ID").WithDataType(FieldTypeInt64).WithIsPrimaryKey(true)). + WithField(NewField().WithName("vector").WithDataType(FieldTypeFloatVector).WithDim(128)), "ID", }, { "dynamic_schema", NewSchema().WithName("dynamic_schema").WithDescription("dynamic_schema desc").WithAutoID(true).WithDynamicFieldEnabled(true). + WithField(NewField().WithName("ID").WithDataType(FieldTypeVarChar).WithMaxLength(256)). WithField(NewField().WithName("$meta").WithIsDynamic(true)), "", },