Skip to content

Commit 662c654

Browse files
SimFGSimFG
authored
Add the context param for the hook interceptor (milvus-io#19405)
Signed-off-by: SimFG <bang.fu@zilliz.com> Signed-off-by: SimFG <bang.fu@zilliz.com>
1 parent 2b58bd5 commit 662c654

3 files changed

Lines changed: 14 additions & 12 deletions

File tree

‎api/hook/hook.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package hook
22

3+
import "context"
4+
35
type Hook interface {
46
Init(params map[string]string) error
5-
Mock(req interface{}, fullMethod string) (bool, interface{}, error)
6-
Before(req interface{}, fullMethod string) error
7-
After(result interface{}, err error, fullMethod string) error
7+
Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error)
8+
Before(ctx context.Context, req interface{}, fullMethod string) error
9+
After(ctx context.Context, result interface{}, err error, fullMethod string) error
810
Release()
911
}

‎internal/proxy/hook_interceptor.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,15 @@ func (d defaultHook) Init(params map[string]string) error {
1818
return nil
1919
}
2020

21-
func (d defaultHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
21+
func (d defaultHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
2222
return false, nil, nil
2323
}
2424

25-
func (d defaultHook) Before(req interface{}, fullMethod string) error {
25+
func (d defaultHook) Before(ctx context.Context, req interface{}, fullMethod string) error {
2626
return nil
2727
}
2828

29-
func (d defaultHook) After(result interface{}, err error, fullMethod string) error {
29+
func (d defaultHook) After(ctx context.Context, result interface{}, err error, fullMethod string) error {
3030
return nil
3131
}
3232

@@ -79,15 +79,15 @@ func UnaryServerHookInterceptor() grpc.UnaryServerInterceptor {
7979
err error
8080
)
8181

82-
if isMock, mockResp, err = hoo.Mock(req, fullMethod); isMock {
82+
if isMock, mockResp, err = hoo.Mock(ctx, req, fullMethod); isMock {
8383
return mockResp, err
8484
}
8585

86-
if err = hoo.Before(req, fullMethod); err != nil {
86+
if err = hoo.Before(ctx, req, fullMethod); err != nil {
8787
return nil, err
8888
}
8989
realResp, realErr = handler(ctx, req)
90-
if err = hoo.After(realResp, realErr, fullMethod); err != nil {
90+
if err = hoo.After(ctx, realResp, realErr, fullMethod); err != nil {
9191
return nil, err
9292
}
9393
return realResp, realErr

‎internal/proxy/hook_interceptor_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ type mockHook struct {
2828
mockErr error
2929
}
3030

31-
func (m mockHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
31+
func (m mockHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
3232
return true, m.mockRes, m.mockErr
3333
}
3434

@@ -42,7 +42,7 @@ type beforeMock struct {
4242
err error
4343
}
4444

45-
func (b beforeMock) Before(r interface{}, fullMethod string) error {
45+
func (b beforeMock) Before(ctx context.Context, r interface{}, fullMethod string) error {
4646
re, ok := r.(*req)
4747
if !ok {
4848
return errors.New("r is invalid type")
@@ -61,7 +61,7 @@ type afterMock struct {
6161
err error
6262
}
6363

64-
func (a afterMock) After(r interface{}, err error, fullMethod string) error {
64+
func (a afterMock) After(ctx context.Context, r interface{}, err error, fullMethod string) error {
6565
re, ok := r.(*resp)
6666
if !ok {
6767
return errors.New("r is invalid type")

0 commit comments

Comments
 (0)