milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add the `context` param for the `hook` interceptor (#19405) 

Signed-off-by: SimFG <bang.fu@zilliz.com>

Signed-off-by: SimFG <bang.fu@zilliz.com>
pull/19388/head
SimFG 2022-09-24 09:56:51 +08:00 committed by GitHub
parent 2b58bd5c0a
commit 662c654fe3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 14 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
api/hook/hook.go
View File

@ -1,9 +1,11 @@
package hook
import "context"
type Hook interface {
Init(params map[string]string) error
Mock(req interface{}, fullMethod string) (bool, interface{}, error)
Before(req interface{}, fullMethod string) error
After(result interface{}, err error, fullMethod string) error
Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error)
Before(ctx context.Context, req interface{}, fullMethod string) error
After(ctx context.Context, result interface{}, err error, fullMethod string) error
Release()
}

12
internal/proxy/hook_interceptor.go
View File

@ -18,15 +18,15 @@ func (d defaultHook) Init(params map[string]string) error {
return nil
}
func (d defaultHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
func (d defaultHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
return false, nil, nil
}
func (d defaultHook) Before(req interface{}, fullMethod string) error {
func (d defaultHook) Before(ctx context.Context, req interface{}, fullMethod string) error {
return nil
}
func (d defaultHook) After(result interface{}, err error, fullMethod string) error {
func (d defaultHook) After(ctx context.Context, result interface{}, err error, fullMethod string) error {
return nil
}
@ -79,15 +79,15 @@ func UnaryServerHookInterceptor() grpc.UnaryServerInterceptor {
err error
)
if isMock, mockResp, err = hoo.Mock(req, fullMethod); isMock {
if isMock, mockResp, err = hoo.Mock(ctx, req, fullMethod); isMock {
return mockResp, err
}
if err = hoo.Before(req, fullMethod); err != nil {
if err = hoo.Before(ctx, req, fullMethod); err != nil {
return nil, err
}
realResp, realErr = handler(ctx, req)
if err = hoo.After(realResp, realErr, fullMethod); err != nil {
if err = hoo.After(ctx, realResp, realErr, fullMethod); err != nil {
return nil, err
}
return realResp, realErr

6
internal/proxy/hook_interceptor_test.go
View File

@ -28,7 +28,7 @@ type mockHook struct {
mockErr error
}
func (m mockHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
func (m mockHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
return true, m.mockRes, m.mockErr
}
@ -42,7 +42,7 @@ type beforeMock struct {
err error
}
func (b beforeMock) Before(r interface{}, fullMethod string) error {
func (b beforeMock) Before(ctx context.Context, r interface{}, fullMethod string) error {
re, ok := r.(*req)
if !ok {
return errors.New("r is invalid type")
@ -61,7 +61,7 @@ type afterMock struct {
err error
}
func (a afterMock) After(r interface{}, err error, fullMethod string) error {
func (a afterMock) After(ctx context.Context, r interface{}, err error, fullMethod string) error {
re, ok := r.(*resp)
if !ok {
return errors.New("r is invalid type")