diff --git a/client/admin.go b/client/admin.go index 2fad21f3d..ec01ed1ef 100644 --- a/client/admin.go +++ b/client/admin.go @@ -15,6 +15,8 @@ import ( "context" server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + + "github.com/milvus-io/milvus-sdk-go/v2/entity" ) // GetVersion returns milvus server version information. @@ -28,3 +30,25 @@ func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) { } return resp.GetVersion(), nil } + +// CheckHealth returns milvus state +func (c *GrpcClient) CheckHealth(ctx context.Context) (*entity.MilvusState, error) { + if c.Service == nil { + return nil, ErrClientNotReady + } + resp, err := c.Service.CheckHealth(ctx, &server.CheckHealthRequest{}) + if err != nil { + return nil, err + } + + states := make([]entity.QuotaState, 0, len(resp.GetQuotaStates())) + for _, state := range resp.GetQuotaStates() { + states = append(states, entity.QuotaState(state)) + } + + return &entity.MilvusState{ + IsHealthy: resp.GetIsHealthy(), + Reasons: resp.GetReasons(), + QuotaStates: states, + }, nil +} diff --git a/client/admin_test.go b/client/admin_test.go new file mode 100644 index 000000000..303c2888d --- /dev/null +++ b/client/admin_test.go @@ -0,0 +1,55 @@ +package client + +import ( + "context" + "github.com/milvus-io/milvus-sdk-go/v2/entity" + "github.com/stretchr/testify/assert" + "testing" + + "github.com/golang/protobuf/proto" + common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" + server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" +) + +func TestCheckHealth(t *testing.T) { + ctx := context.Background() + c := testClient(ctx, t) + defer c.Close() + + t.Run("test milvus healthy", func(t *testing.T) { + mockServer.SetInjection(MCheckHealth, func(ctx context.Context, message proto.Message) (proto.Message, error) { + resp := &server.CheckHealthResponse{ + Status: &common.Status{ErrorCode: common.ErrorCode_Success}, + IsHealthy: true, + Reasons: nil, + QuotaStates: nil, + } + return resp, nil + }) + + resp, err := c.CheckHealth(ctx) + assert.Nil(t, err) + assert.True(t, resp.IsHealthy) + assert.Empty(t, resp.Reasons) + assert.Empty(t, resp.QuotaStates) + }) + + t.Run("test milvus unhealthy", func(t *testing.T) { + mockServer.SetInjection(MCheckHealth, func(ctx context.Context, message proto.Message) (proto.Message, error) { + resp := &server.CheckHealthResponse{ + Status: &common.Status{ErrorCode: common.ErrorCode_Success}, + IsHealthy: false, + Reasons: []string{"some reason"}, + QuotaStates: []server.QuotaState{server.QuotaState_DenyToRead, server.QuotaState_DenyToWrite}, + } + return resp, nil + }) + + resp, err := c.CheckHealth(ctx) + assert.Nil(t, err) + assert.False(t, resp.IsHealthy) + assert.ElementsMatch(t, resp.Reasons, []string{"some reason"}) + assert.ElementsMatch(t, resp.QuotaStates, []entity.QuotaState{entity.QuotaStateDenyToRead, entity.QuotaStateDenyToWrite}) + }) + +} diff --git a/client/client.go b/client/client.go index fc44d6f22..74c545243 100644 --- a/client/client.go +++ b/client/client.go @@ -205,6 +205,8 @@ type Client interface { // GetVersion get milvus version GetVersion(ctx context.Context) (string, error) + // CheckHealth returns milvus state + CheckHealth(ctx context.Context) (*entity.MilvusState, error) } // NewClient create a client connected to remote milvus cluster. diff --git a/entity/check_health.go b/entity/check_health.go new file mode 100644 index 000000000..c821d8c9e --- /dev/null +++ b/entity/check_health.go @@ -0,0 +1,22 @@ +package entity + +type QuotaState int32 + +const ( + // QuotaStateUnknown zero value placeholder + QuotaStateUnknown QuotaState = 0 + // QuotaStateReadLimited too many read tasks, read requests are limited + QuotaStateReadLimited QuotaState = 2 + // QuotaStateWriteLimited too many write tasks, write requests are limited + QuotaStateWriteLimited QuotaState = 3 + // QuotaStateDenyToRead too many read tasks, temporarily unable to process read requests + QuotaStateDenyToRead QuotaState = 4 + // QuotaStateDenyToWrite too many write tasks, temporarily unable to process write requests + QuotaStateDenyToWrite QuotaState = 5 +) + +type MilvusState struct { + IsHealthy bool + Reasons []string + QuotaStates []QuotaState +}