enhance: [GoSDK] Sync API names and add missing APIs

Related to milvus-io#31293

- Rename `UsingDatabase` to `UseDatabase`
- Uncomment default value methods
- Add missing RBAC APIs
- Add some resource group APIs

Signed-off-by: Congqi Xia <[email protected]>

client/entity/common.go

@@ -16,6 +16,8 @@
 
 package entity
 
+import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+
 // MetricType metric type
 type MetricType string
 
@@ -31,3 +33,12 @@ const (
 	SUPERSTRUCTURE MetricType = "SUPERSTRUCTURE"
 	BM25           MetricType = "BM25"
 )
+
+// CompactionState enum type for compaction state
+type CompactionState commonpb.CompactionState
+
+// CompactionState Constants
+const (
+	CompactionStateRunning   CompactionState = CompactionState(commonpb.CompactionState_Executing)
+	CompactionStateCompleted CompactionState = CompactionState(commonpb.CompactionState_Completed)
+)

client/entity/field.go

@@ -193,6 +193,8 @@ type Field struct {
 	IsPartitionKey  bool
 	IsClusteringKey bool
 	ElementType     FieldType
+	DefaultValue    *schemapb.ValueField
+	Nullable        bool
 }
 
 // ProtoMessage generates corresponding FieldSchema
@@ -261,7 +263,11 @@ func (f *Field) WithIsClusteringKey(isClusteringKey bool) *Field {
 	return f
 }
 
-/*
+func (f *Field) WithNullable(nullable bool) *Field {
+	f.Nullable = nullable
+	return f
+}
+
 func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
 	f.DefaultValue = &schemapb.ValueField{
 		Data: &schemapb.ValueField_BoolData{
@@ -314,7 +320,7 @@ func (f *Field) WithDefaultValueString(defaultValue string) *Field {
 		},
 	}
 	return f
-}*/
+}
 
 func (f *Field) WithTypeParams(key string, value string) *Field {
 	if f.TypeParams == nil {

client/entity/field_test.go

@@ -30,13 +30,13 @@ func TestFieldSchema(t *testing.T) {
 		NewField().WithName("array_field").WithDataType(FieldTypeArray).WithElementType(FieldTypeBool).WithMaxCapacity(128),
 		NewField().WithName("clustering_key").WithDataType(FieldTypeInt32).WithIsClusteringKey(true),
 		NewField().WithName("varchar_text").WithDataType(FieldTypeVarChar).WithMaxLength(65535).WithEnableAnalyzer(true).WithAnalyzerParams(map[string]any{}),
-		/*
-			NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
-			NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
-			NewField().WithName("default_value_long").WithDataType(FieldTypeInt64).WithDefaultValueLong(1),
-			NewField().WithName("default_value_float").WithDataType(FieldTypeFloat).WithDefaultValueFloat(1),
-			NewField().WithName("default_value_double").WithDataType(FieldTypeDouble).WithDefaultValueDouble(1),
-			NewField().WithName("default_value_string").WithDataType(FieldTypeString).WithDefaultValueString("a"),*/
+
+		NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
+		NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
+		NewField().WithName("default_value_long").WithDataType(FieldTypeInt64).WithDefaultValueLong(1),
+		NewField().WithName("default_value_float").WithDataType(FieldTypeFloat).WithDefaultValueFloat(1),
+		NewField().WithName("default_value_double").WithDataType(FieldTypeDouble).WithDefaultValueDouble(1),
+		NewField().WithName("default_value_string").WithDataType(FieldTypeString).WithDefaultValueString("a"),
 	}
 
 	for _, field := range fields {

client/entity/load_state.go

@@ -0,0 +1,34 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 entity
+
+import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+
+type LoadStateCode commonpb.LoadState
+
+const (
+	// LoadStateNone      LoadStateCode = LoadStateCode(commonpb.LoadState)
+	LoadStateLoading   LoadStateCode = LoadStateCode(commonpb.LoadState_LoadStateLoading)
+	LoadStateLoaded    LoadStateCode = LoadStateCode(commonpb.LoadState_LoadStateLoaded)
+	LoadStateUnloading LoadStateCode = LoadStateCode(commonpb.LoadState_LoadStateNotExist)
+	LoadStateNotLoad   LoadStateCode = LoadStateCode(commonpb.LoadState_LoadStateNotLoad)
+)
+
+type LoadState struct {
+	State    LoadStateCode
+	Progress int64
+}

client/entity/rbac.go

@@ -0,0 +1,35 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 entity
+
+type User struct {
+	UserName string
+	Roles    []string
+}
+
+type Role struct {
+	RoleName   string
+	Privileges []GrantItem
+}
+
+type GrantItem struct {
+	Object     string
+	ObjectName string
+	RoleName   string
+	Grantor    string
+	Privilege  string
+}

client/milvusclient/collection.go

@@ -147,3 +147,23 @@ func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOpti
 		return merr.CheckRPCCall(resp, err)
 	})
 }
+
+type GetCollectionOption interface {
+	Request() *milvuspb.GetCollectionStatisticsRequest
+}
+
+func (c *Client) GetCollectionStats(ctx context.Context, opt GetCollectionOption) (map[string]string, error) {
+	var stats map[string]string
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.GetCollectionStatistics(ctx, opt.Request())
+		if err = merr.CheckRPCCall(resp, err); err != nil {
+			return err
+		}
+		stats = entity.KvPairsMap(resp.GetStats())
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return stats, nil
+}

client/milvusclient/collection_options.go

@@ -310,3 +310,17 @@ func (opt *alterCollectionOption) Request() *milvuspb.AlterCollectionRequest {
 func NewAlterCollectionOption(collection string) *alterCollectionOption {
 	return &alterCollectionOption{collectionName: collection, properties: make(map[string]string)}
 }
+
+type getCollectionStatsOption struct {
+	collectionName string
+}
+
+func (opt *getCollectionStatsOption) Request() *milvuspb.GetCollectionStatisticsRequest {
+	return &milvuspb.GetCollectionStatisticsRequest{
+		CollectionName: opt.collectionName,
+	}
+}
+
+func NewGetCollectionStatsOption(collectionName string) *getCollectionStatsOption {
+	return &getCollectionStatsOption{collectionName: collectionName}
+}

client/milvusclient/collection_test.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
@@ -315,6 +316,38 @@ func (s *CollectionSuite) TestAlterCollection() {
 	})
 }
 
+func (s *CollectionSuite) TestGetCollectionStats() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	s.Run("success", func() {
+		collName := fmt.Sprintf("coll_%s", s.randString(6))
+		s.mock.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, gcsr *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
+			s.Equal(collName, gcsr.GetCollectionName())
+			return &milvuspb.GetCollectionStatisticsResponse{
+				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+				Stats: []*commonpb.KeyValuePair{
+					{Key: "row_count", Value: "1000"},
+				},
+			}, nil
+		}).Once()
+
+		stats, err := s.client.GetCollectionStats(ctx, NewGetCollectionStatsOption(collName))
+		s.NoError(err)
+
+		s.Len(stats, 1)
+		s.Equal("1000", stats["row_count"])
+	})
+
+	s.Run("failure", func() {
+		collName := fmt.Sprintf("coll_%s", s.randString(6))
+		s.mock.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(nil, errors.New("mocked")).Once()
+
+		_, err := s.client.GetCollectionStats(ctx, NewGetCollectionStatsOption(collName))
+		s.Error(err)
+	})
+}
+
 func TestCollection(t *testing.T) {
 	suite.Run(t, new(CollectionSuite))
 }

client/milvusclient/database.go

@@ -25,7 +25,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/merr"
 )
 
-func (c *Client) UsingDatabase(ctx context.Context, option UsingDatabaseOption) error {
+func (c *Client) UseDatabase(ctx context.Context, option UseDatabaseOption) error {
 	dbName := option.DbName()
 	c.usingDatabase(dbName)
 	return c.connectInternal(ctx)

client/milvusclient/database_options.go

@@ -18,20 +18,20 @@ package milvusclient
 
 import "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 
-type UsingDatabaseOption interface {
+type UseDatabaseOption interface {
 	DbName() string
 }
 
-type usingDatabaseNameOpt struct {
+type useDatabaseNameOpt struct {
 	dbName string
 }
 
-func (opt *usingDatabaseNameOpt) DbName() string {
+func (opt *useDatabaseNameOpt) DbName() string {
 	return opt.dbName
 }
 
-func NewUsingDatabaseOption(dbName string) *usingDatabaseNameOpt {
-	return &usingDatabaseNameOpt{
+func NewUseDatabaseOption(dbName string) *useDatabaseNameOpt {
+	return &useDatabaseNameOpt{
 		dbName: dbName,
 	}
 }

client/milvusclient/database_test.go

@@ -88,6 +88,26 @@ func (s *DatabaseSuite) TestDropDatabase() {
 	})
 }
 
+func (s *DatabaseSuite) TestUseDatabase() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	s.Run("success", func() {
+		dbName := fmt.Sprintf("dt_%s", s.randString(6))
+		s.mock.EXPECT().Connect(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cr *milvuspb.ConnectRequest) (*milvuspb.ConnectResponse, error) {
+			return &milvuspb.ConnectResponse{
+				Status:     merr.Success(),
+				ServerInfo: &commonpb.ServerInfo{},
+			}, nil
+		}).Once()
+
+		err := s.client.UseDatabase(ctx, NewUseDatabaseOption(dbName))
+		s.NoError(err)
+
+		s.Equal(dbName, s.client.currentDB)
+	})
+}
+
 func TestDatabase(t *testing.T) {
 	suite.Run(t, new(DatabaseSuite))
 }