Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Dec 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions client/entity/common.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

package entity

import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"

// MetricType metric type
type MetricType string

Expand All @@ -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)
)
10 changes: 8 additions & 2 deletions client/entity/field.go
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,8 @@
IsPartitionKey bool
IsClusteringKey bool
ElementType FieldType
DefaultValue *schemapb.ValueField
Nullable bool
}

// ProtoMessage generates corresponding FieldSchema
Expand Down Expand Up @@ -261,7 +263,11 @@
return f
}

/*
func (f *Field) WithNullable(nullable bool) *Field {
f.Nullable = nullable
return f

Check warning on line 268 in client/entity/field.go

View check run for this annotation

Codecov / codecov/patch

client/entity/field.go#L266-L268

Added lines #L266 - L268 were not covered by tests
}

func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
f.DefaultValue = &schemapb.ValueField{
Data: &schemapb.ValueField_BoolData{
Expand Down Expand Up @@ -314,7 +320,7 @@
},
}
return f
}*/
}

func (f *Field) WithTypeParams(key string, value string) *Field {
if f.TypeParams == nil {
Expand Down
14 changes: 7 additions & 7 deletions client/entity/field_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
34 changes: 34 additions & 0 deletions client/entity/load_state.go
Original file line number Diff line number Diff line change
@@ -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
}
35 changes: 35 additions & 0 deletions client/entity/rbac.go
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 1 addition & 1 deletion client/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ require (
github.com/blang/semver/v4 v4.0.0
github.com/cockroachdb/errors v1.9.1
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69
github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b
github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84
github.com/quasilyte/go-ruleguard/dsl v0.3.22
github.com/samber/lo v1.27.0
Expand Down
4 changes: 2 additions & 2 deletions client/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -318,8 +318,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr
github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69 h1:Qt0Bv2Fum3EX3OlkuQYHJINBzeU4oEuHy2lXSfB/gZw=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b h1:iPPhnFx+s7FF53UeWj7A4EYhPRMFPL6mHqyQw7qRjeQ=
github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 h1:EAFxmxUVp5yYFDCrX1MQoSxkTO+ycy8NXEqEDEB3cRM=
github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84/go.mod h1:RATa0GS4jhkPpsYOvQ/QvcNz8rd+TlRPDiSyXQnMMxs=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
Expand Down
20 changes: 20 additions & 0 deletions client/milvusclient/collection.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
14 changes: 14 additions & 0 deletions client/milvusclient/collection_options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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}
}
33 changes: 33 additions & 0 deletions client/milvusclient/collection_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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))
}
2 changes: 1 addition & 1 deletion client/milvusclient/database.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 5 additions & 5 deletions client/milvusclient/database_options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}
Expand Down
20 changes: 20 additions & 0 deletions client/milvusclient/database_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
Loading
Loading