Add the context param for the hook interceptor (milvus-io#19405)

Signed-off-by: SimFG <[email protected]> Signed-off-by: SimFG <[email protected]>
sunby · Sep 24, 2022 · 662c654 · 662c654
1 parent 2b58bd5
commit 662c654
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 12 deletions.
diff --git a/api/hook/hook.go b/api/hook/hook.go
@@ -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()
 }
diff --git a/internal/proxy/hook_interceptor.go b/internal/proxy/hook_interceptor.go
@@ -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

diff --git a/internal/proxy/hook_interceptor_test.go b/internal/proxy/hook_interceptor_test.go
@@ -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")