enhance: Support to trace restful request and request error (#28685)

issue: #28348

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

internal/distributed/proxy/httpserver/handler.go

@@ -1,6 +1,7 @@
 package httpserver
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/gin-gonic/gin"
@@ -10,9 +11,12 @@ import (
 	"github.com/milvus-io/milvus/internal/types"
 )
 
+type RestRequestInterceptor func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error)
+
 // Handlers handles http requests
 type Handlers struct {
-	proxy types.ProxyComponent
+	proxy        types.ProxyComponent
+	interceptors []RestRequestInterceptor
 }
 
 // NewHandlers creates a new Handlers

internal/distributed/proxy/httpserver/handler_v1.go

@@ -6,11 +6,13 @@ import (
 	"net/http"
 	"strconv"
 
+	"github.com/cockroachdb/errors"
 	"github.com/gin-gonic/gin"
 	"github.com/gin-gonic/gin/binding"
 	"github.com/golang/protobuf/proto"
 	"github.com/tidwall/gjson"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
@@ -19,27 +21,30 @@ import (
 	"github.com/milvus-io/milvus/pkg/common"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/merr"
+	"github.com/milvus-io/milvus/pkg/util/requestutil"
 )
 
+var RestRequestInterceptorErr = errors.New("interceptor error placeholder")
+
 func checkAuthorization(ctx context.Context, c *gin.Context, req interface{}) error {
 	if proxy.Params.CommonCfg.AuthorizationEnabled.GetAsBool() {
 		username, ok := c.Get(ContextUsername)
 		if !ok || username.(string) == "" {
 			c.JSON(http.StatusUnauthorized, gin.H{HTTPReturnCode: merr.Code(merr.ErrNeedAuthenticate), HTTPReturnMessage: merr.ErrNeedAuthenticate.Error()})
-			return merr.ErrNeedAuthenticate
+			return RestRequestInterceptorErr
 		}
 		_, authErr := proxy.PrivilegeInterceptor(ctx, req)
 		if authErr != nil {
 			c.JSON(http.StatusForbidden, gin.H{HTTPReturnCode: merr.Code(authErr), HTTPReturnMessage: authErr.Error()})
-			return authErr
+			return RestRequestInterceptorErr
 		}
 	}
 	return nil
 }
 
-func (h *Handlers) checkDatabase(ctx context.Context, c *gin.Context, dbName string) bool {
+func (h *Handlers) checkDatabase(ctx context.Context, c *gin.Context, dbName string) error {
 	if dbName == DefaultDbName {
-		return true
+		return nil
 	}
 	response, err := h.proxy.ListDatabases(ctx, &milvuspb.ListDatabasesRequest{})
 	if err == nil {
@@ -47,30 +52,25 @@ func (h *Handlers) checkDatabase(ctx context.Context, c *gin.Context, dbName str
 	}
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
-		return false
+		return RestRequestInterceptorErr
 	}
 	for _, db := range response.DbNames {
 		if db == dbName {
-			return true
+			return nil
 		}
 	}
 	c.AbortWithStatusJSON(http.StatusOK, gin.H{
 		HTTPReturnCode:    merr.Code(merr.ErrDatabaseNotFound),
 		HTTPReturnMessage: merr.ErrDatabaseNotFound.Error() + ", database: " + dbName,
 	})
-	return false
+	return RestRequestInterceptorErr
 }
 
-func (h *Handlers) describeCollection(ctx context.Context, c *gin.Context, dbName string, collectionName string, needAuth bool) (*milvuspb.DescribeCollectionResponse, error) {
+func (h *Handlers) describeCollection(ctx context.Context, c *gin.Context, dbName string, collectionName string) (*milvuspb.DescribeCollectionResponse, error) {
 	req := milvuspb.DescribeCollectionRequest{
 		DbName:         dbName,
 		CollectionName: collectionName,
 	}
-	if needAuth {
-		if err := checkAuthorization(ctx, c, &req); err != nil {
-			return nil, err
-		}
-	}
 	response, err := h.proxy.DescribeCollection(ctx, &req)
 	if err == nil {
 		err = merr.Error(response.GetStatus())
@@ -104,6 +104,7 @@ func (h *Handlers) hasCollection(ctx context.Context, c *gin.Context, dbName str
 }
 
 func (h *Handlers) RegisterRoutesToV1(router gin.IRouter) {
+	h.registerRestRequestInterceptor()
 	router.GET(VectorCollectionsPath, h.listCollections)
 	router.POST(VectorCollectionsCreatePath, h.createCollection)
 	router.GET(VectorCollectionsDescribePath, h.getCollectionDetails)
@@ -116,27 +117,74 @@ func (h *Handlers) RegisterRoutesToV1(router gin.IRouter) {
 	router.POST(VectorSearchPath, h.search)
 }
 
+func (h *Handlers) registerRestRequestInterceptor() {
+	h.interceptors = []RestRequestInterceptor{
+		// authorization
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			err := checkAuthorization(ctx, ginCtx, req)
+			if err != nil {
+				return nil, err
+			}
+			return handler(ctx, req)
+		},
+		// check database
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			value, ok := requestutil.GetDbNameFromRequest(req)
+			if !ok {
+				return handler(ctx, req)
+			}
+			err := h.checkDatabase(ctx, ginCtx, value.(string))
+			if err != nil {
+				return nil, err
+			}
+			return handler(ctx, req)
+		},
+		// trace request
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			return proxy.TraceLogInterceptor(ctx, req, &grpc.UnaryServerInfo{
+				FullMethod: ginCtx.Request.URL.Path,
+			}, handler)
+		},
+	}
+}
+
+func (h *Handlers) executeRestRequestInterceptor(ctx context.Context,
+	ginCtx *gin.Context,
+	req any, handler func(reqCtx context.Context, req any) (any, error),
+) (any, error) {
+	f := handler
+	for i := len(h.interceptors) - 1; i >= 0; i-- {
+		f = func(j int, handlerFunc func(reqCtx context.Context, req any) (any, error)) func(reqCtx context.Context, req any) (any, error) {
+			return func(reqCtx context.Context, req any) (any, error) {
+				return h.interceptors[j](reqCtx, ginCtx, req, handlerFunc)
+			}
+		}(i, f)
+	}
+	return f(ctx, req)
+}
+
 func (h *Handlers) listCollections(c *gin.Context) {
 	dbName := c.DefaultQuery(HTTPDbName, DefaultDbName)
-	req := milvuspb.ShowCollectionsRequest{
+	req := &milvuspb.ShowCollectionsRequest{
 		DbName: dbName,
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+
+	resp, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		return h.proxy.ShowCollections(reqCtx, req.(*milvuspb.ShowCollectionsRequest))
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, dbName) {
-		return
-	}
-	response, err := h.proxy.ShowCollections(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(resp.(*milvuspb.ShowCollectionsResponse).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 		return
 	}
+	response := resp.(*milvuspb.ShowCollectionsResponse)
 	var collections []string
 	if response.CollectionNames != nil {
 		collections = response.CollectionNames
@@ -204,7 +252,7 @@ func (h *Handlers) createCollection(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.CreateCollectionRequest{
+	req := &milvuspb.CreateCollectionRequest{
 		DbName:           httpReq.DbName,
 		CollectionName:   httpReq.CollectionName,
 		Schema:           schema,
@@ -213,22 +261,21 @@ func (h *Handlers) createCollection(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		return h.proxy.CreateCollection(reqCtx, req.(*milvuspb.CreateCollectionRequest))
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	response, err := h.proxy.CreateCollection(ctx, &req)
 	if err == nil {
-		err = merr.Error(response)
+		err = merr.Error(response.(*commonpb.Status))
 	}
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 		return
 	}
 
-	response, err = h.proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{
+	statusResponse, err := h.proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 		FieldName:      httpReq.VectorField,
@@ -236,18 +283,18 @@ func (h *Handlers) createCollection(c *gin.Context) {
 		ExtraParams:    []*commonpb.KeyValuePair{{Key: common.MetricTypeKey, Value: httpReq.MetricType}},
 	})
 	if err == nil {
-		err = merr.Error(response)
+		err = merr.Error(statusResponse)
 	}
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 		return
 	}
-	response, err = h.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
+	statusResponse, err = h.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 	})
 	if err == nil {
-		err = merr.Error(response)
+		err = merr.Error(statusResponse)
 	}
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
@@ -269,13 +316,30 @@ func (h *Handlers) getCollectionDetails(c *gin.Context) {
 	dbName := c.DefaultQuery(HTTPDbName, DefaultDbName)
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), dbName)
-	if !h.checkDatabase(ctx, c, dbName) {
-		return
+
+	req := &milvuspb.DescribeCollectionRequest{
+		DbName:         dbName,
+		CollectionName: collectionName,
+	}
+
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		return h.proxy.DescribeCollection(reqCtx, req.(*milvuspb.DescribeCollectionRequest))
+	})
+
+	if err == nil {
+		err = merr.Error(response.(*milvuspb.DescribeCollectionResponse).GetStatus())
 	}
-	coll, err := h.describeCollection(ctx, c, dbName, collectionName, true)
 	if err != nil {
+		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 		return
 	}
+	coll := response.(*milvuspb.DescribeCollectionResponse)
+	primaryField, ok := getPrimaryField(coll.Schema)
+	if ok && primaryField.AutoID && !coll.Schema.AutoID {
+		log.Warn("primary filed autoID VS schema autoID", zap.String("collectionName", collectionName), zap.Bool("primary Field", primaryField.AutoID), zap.Bool("schema", coll.Schema.AutoID))
+		coll.Schema.AutoID = EnableAutoID
+	}
+
 	stateResp, err := h.proxy.GetLoadState(ctx, &milvuspb.GetLoadStateRequest{
 		DbName:         dbName,
 		CollectionName: collectionName,
@@ -349,32 +413,31 @@ func (h *Handlers) dropCollection(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.DropCollectionRequest{
+	req := &milvuspb.DropCollectionRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		has, err := h.hasCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
+		if err != nil {
+			return nil, RestRequestInterceptorErr
+		}
+		if !has {
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrCollectionNotFound),
+				HTTPReturnMessage: merr.ErrCollectionNotFound.Error() + ", database: " + httpReq.DbName + ", collection: " + httpReq.CollectionName,
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		return h.proxy.DropCollection(reqCtx, req.(*milvuspb.DropCollectionRequest))
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	has, err := h.hasCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
-	if err != nil {
-		return
-	}
-	if !has {
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrCollectionNotFound),
-			HTTPReturnMessage: merr.ErrCollectionNotFound.Error() + ", database: " + httpReq.DbName + ", collection: " + httpReq.CollectionName,
-		})
-		return
-	}
-	response, err := h.proxy.DropCollection(ctx, &req)
 	if err == nil {
-		err = merr.Error(response)
+		err = merr.Error(response.(*commonpb.Status))
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
@@ -405,7 +468,7 @@ func (h *Handlers) query(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.QueryRequest{
+	req := &milvuspb.QueryRequest{
 		DbName:             httpReq.DbName,
 		CollectionName:     httpReq.CollectionName,
 		Expr:               httpReq.Filter,
@@ -421,21 +484,21 @@ func (h *Handlers) query(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		return h.proxy.Query(reqCtx, req.(*milvuspb.QueryRequest))
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	response, err := h.proxy.Query(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.QueryResults).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 	} else {
+		queryResp := response.(*milvuspb.QueryResults)
 		allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
-		outputData, err := buildQueryResp(int64(0), response.OutputFields, response.FieldsData, nil, nil, allowJS)
+		outputData, err := buildQueryResp(int64(0), queryResp.OutputFields, queryResp.FieldsData, nil, nil, allowJS)
 		if err != nil {
 			log.Warn("high level restful api, fail to deal with query result", zap.Any("response", response), zap.Error(err))
 			c.JSON(http.StatusOK, gin.H{
@@ -469,7 +532,7 @@ func (h *Handlers) get(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.QueryRequest{
+	req := &milvuspb.QueryRequest{
 		DbName:             httpReq.DbName,
 		CollectionName:     httpReq.CollectionName,
 		OutputFields:       httpReq.OutputFields,
@@ -477,35 +540,36 @@ func (h *Handlers) get(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
+		if err != nil || coll == nil {
+			return nil, RestRequestInterceptorErr
+		}
+		body, _ := c.Get(gin.BodyBytesKey)
+		filter, err := checkGetPrimaryKey(coll.Schema, gjson.Get(string(body.([]byte)), DefaultPrimaryFieldName))
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
+				HTTPReturnMessage: merr.ErrCheckPrimaryKey.Error() + ", error: " + err.Error(),
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		queryReq := req.(*milvuspb.QueryRequest)
+		queryReq.Expr = filter
+		return h.proxy.Query(reqCtx, queryReq)
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName, false)
-	if err != nil || coll == nil {
-		return
-	}
-	body, _ := c.Get(gin.BodyBytesKey)
-	filter, err := checkGetPrimaryKey(coll.Schema, gjson.Get(string(body.([]byte)), DefaultPrimaryFieldName))
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
-			HTTPReturnMessage: merr.ErrCheckPrimaryKey.Error() + ", error: " + err.Error(),
-		})
-		return
-	}
-	req.Expr = filter
-	response, err := h.proxy.Query(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.QueryResults).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 	} else {
+		queryResp := response.(*milvuspb.QueryResults)
 		allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
-		outputData, err := buildQueryResp(int64(0), response.OutputFields, response.FieldsData, nil, nil, allowJS)
+		outputData, err := buildQueryResp(int64(0), queryResp.OutputFields, queryResp.FieldsData, nil, nil, allowJS)
 		if err != nil {
 			log.Warn("high level restful api, fail to deal with get result", zap.Any("response", response), zap.Error(err))
 			c.JSON(http.StatusOK, gin.H{
@@ -538,38 +602,38 @@ func (h *Handlers) delete(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.DeleteRequest{
+	req := &milvuspb.DeleteRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
-		return
-	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName, false)
-	if err != nil || coll == nil {
-		return
-	}
-	req.Expr = httpReq.Filter
-	if req.Expr == "" {
-		body, _ := c.Get(gin.BodyBytesKey)
-		filter, err := checkGetPrimaryKey(coll.Schema, gjson.Get(string(body.([]byte)), DefaultPrimaryFieldName))
-		if err != nil {
-			c.JSON(http.StatusOK, gin.H{
-				HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
-				HTTPReturnMessage: merr.ErrCheckPrimaryKey.Error() + ", error: " + err.Error(),
-			})
-			return
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
+		if err != nil || coll == nil {
+			return nil, RestRequestInterceptorErr
 		}
-		req.Expr = filter
+		deleteReq := req.(*milvuspb.DeleteRequest)
+		deleteReq.Expr = httpReq.Filter
+		if deleteReq.Expr == "" {
+			body, _ := c.Get(gin.BodyBytesKey)
+			filter, err := checkGetPrimaryKey(coll.Schema, gjson.Get(string(body.([]byte)), DefaultPrimaryFieldName))
+			if err != nil {
+				c.JSON(http.StatusOK, gin.H{
+					HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
+					HTTPReturnMessage: merr.ErrCheckPrimaryKey.Error() + ", error: " + err.Error(),
+				})
+				return nil, RestRequestInterceptorErr
+			}
+			deleteReq.Expr = filter
+		}
+		return h.proxy.Delete(ctx, deleteReq)
+	})
+	if err == RestRequestInterceptorErr {
+		return
 	}
-	response, err := h.proxy.Delete(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.MutationResult).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
@@ -606,7 +670,7 @@ func (h *Handlers) insert(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.InsertRequest{
+	req := &milvuspb.InsertRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 		PartitionName:  "_default",
@@ -614,52 +678,53 @@ func (h *Handlers) insert(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
+		if err != nil || coll == nil {
+			return nil, RestRequestInterceptorErr
+		}
+		body, _ := c.Get(gin.BodyBytesKey)
+		err, httpReq.Data = checkAndSetData(string(body.([]byte)), coll)
+		if err != nil {
+			log.Warn("high level restful api, fail to deal with insert data", zap.Any("body", body), zap.Error(err))
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
+				HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		insertReq := req.(*milvuspb.InsertRequest)
+		insertReq.FieldsData, err = anyToColumns(httpReq.Data, coll.Schema)
+		if err != nil {
+			log.Warn("high level restful api, fail to deal with insert data", zap.Any("data", httpReq.Data), zap.Error(err))
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
+				HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		return h.proxy.Insert(ctx, insertReq)
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName, false)
-	if err != nil || coll == nil {
-		return
-	}
-	body, _ := c.Get(gin.BodyBytesKey)
-	err, httpReq.Data = checkAndSetData(string(body.([]byte)), coll)
-	if err != nil {
-		log.Warn("high level restful api, fail to deal with insert data", zap.Any("body", body), zap.Error(err))
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
-			HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
-		})
-		return
-	}
-	req.FieldsData, err = anyToColumns(httpReq.Data, coll.Schema)
-	if err != nil {
-		log.Warn("high level restful api, fail to deal with insert data", zap.Any("data", httpReq.Data), zap.Error(err))
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
-			HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
-		})
-		return
-	}
-	response, err := h.proxy.Insert(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.MutationResult).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 	} else {
-		switch response.IDs.GetIdField().(type) {
+		insertResp := response.(*milvuspb.MutationResult)
+		switch insertResp.IDs.GetIdField().(type) {
 		case *schemapb.IDs_IntId:
 			allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
 			if allowJS {
-				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": response.InsertCnt, "insertIds": response.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
+				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
 			} else {
-				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": response.InsertCnt, "insertIds": formatInt64(response.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
+				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": formatInt64(insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
 			}
 		case *schemapb.IDs_StrId:
-			c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": response.InsertCnt, "insertIds": response.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
+			c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
 		default:
 			c.JSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
@@ -697,7 +762,7 @@ func (h *Handlers) upsert(c *gin.Context) {
 		})
 		return
 	}
-	req := milvuspb.UpsertRequest{
+	req := &milvuspb.UpsertRequest{
 		DbName:         httpReq.DbName,
 		CollectionName: httpReq.CollectionName,
 		PartitionName:  "_default",
@@ -705,57 +770,58 @@ func (h *Handlers) upsert(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName)
+		if err != nil || coll == nil {
+			return nil, RestRequestInterceptorErr
+		}
+		if coll.Schema.AutoID {
+			err := merr.WrapErrParameterInvalid("autoID: false", "autoID: true", "cannot upsert an autoID collection")
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
+			return nil, RestRequestInterceptorErr
+		}
+		body, _ := c.Get(gin.BodyBytesKey)
+		err, httpReq.Data = checkAndSetData(string(body.([]byte)), coll)
+		if err != nil {
+			log.Warn("high level restful api, fail to deal with upsert data", zap.Any("body", body), zap.Error(err))
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
+				HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		upsertReq := req.(*milvuspb.UpsertRequest)
+		upsertReq.FieldsData, err = anyToColumns(httpReq.Data, coll.Schema)
+		if err != nil {
+			log.Warn("high level restful api, fail to deal with upsert data", zap.Any("data", httpReq.Data), zap.Error(err))
+			c.AbortWithStatusJSON(http.StatusOK, gin.H{
+				HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
+				HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
+			})
+			return nil, RestRequestInterceptorErr
+		}
+		return h.proxy.Upsert(ctx, upsertReq)
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	coll, err := h.describeCollection(ctx, c, httpReq.DbName, httpReq.CollectionName, false)
-	if err != nil || coll == nil {
-		return
-	}
-	if coll.Schema.AutoID {
-		err := merr.WrapErrParameterInvalid("autoID: false", "autoID: true", "cannot upsert an autoID collection")
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
-		return
-	}
-	body, _ := c.Get(gin.BodyBytesKey)
-	err, httpReq.Data = checkAndSetData(string(body.([]byte)), coll)
-	if err != nil {
-		log.Warn("high level restful api, fail to deal with upsert data", zap.Any("body", body), zap.Error(err))
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
-			HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
-		})
-		return
-	}
-	req.FieldsData, err = anyToColumns(httpReq.Data, coll.Schema)
-	if err != nil {
-		log.Warn("high level restful api, fail to deal with upsert data", zap.Any("data", httpReq.Data), zap.Error(err))
-		c.AbortWithStatusJSON(http.StatusOK, gin.H{
-			HTTPReturnCode:    merr.Code(merr.ErrInvalidInsertData),
-			HTTPReturnMessage: merr.ErrInvalidInsertData.Error() + ", error: " + err.Error(),
-		})
-		return
-	}
-	response, err := h.proxy.Upsert(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.MutationResult).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 	} else {
-		switch response.IDs.GetIdField().(type) {
+		upsertResp := response.(*milvuspb.MutationResult)
+		switch upsertResp.IDs.GetIdField().(type) {
 		case *schemapb.IDs_IntId:
 			allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
 			if allowJS {
-				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": response.UpsertCnt, "upsertIds": response.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
+				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
 			} else {
-				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": response.UpsertCnt, "upsertIds": formatInt64(response.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
+				c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": formatInt64(upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
 			}
 		case *schemapb.IDs_StrId:
-			c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": response.UpsertCnt, "upsertIds": response.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
+			c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
 		default:
 			c.JSON(http.StatusOK, gin.H{
 				HTTPReturnCode:    merr.Code(merr.ErrCheckPrimaryKey),
@@ -796,7 +862,7 @@ func (h *Handlers) search(c *gin.Context) {
 		{Key: ParamRoundDecimal, Value: "-1"},
 		{Key: ParamOffset, Value: strconv.FormatInt(int64(httpReq.Offset), 10)},
 	}
-	req := milvuspb.SearchRequest{
+	req := &milvuspb.SearchRequest{
 		DbName:             httpReq.DbName,
 		CollectionName:     httpReq.CollectionName,
 		Dsl:                httpReq.Filter,
@@ -809,26 +875,26 @@ func (h *Handlers) search(c *gin.Context) {
 	}
 	username, _ := c.Get(ContextUsername)
 	ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
-	if err := checkAuthorization(ctx, c, &req); err != nil {
+	response, err := h.executeRestRequestInterceptor(ctx, c, req, func(reqCtx context.Context, req any) (any, error) {
+		return h.proxy.Search(ctx, req.(*milvuspb.SearchRequest))
+	})
+	if err == RestRequestInterceptorErr {
 		return
 	}
-	if !h.checkDatabase(ctx, c, req.DbName) {
-		return
-	}
-	response, err := h.proxy.Search(ctx, &req)
 	if err == nil {
-		err = merr.Error(response.GetStatus())
+		err = merr.Error(response.(*milvuspb.SearchResults).GetStatus())
 	}
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{HTTPReturnCode: merr.Code(err), HTTPReturnMessage: err.Error()})
 	} else {
-		if response.Results.TopK == int64(0) {
+		searchResp := response.(*milvuspb.SearchResults)
+		if searchResp.Results.TopK == int64(0) {
 			c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: []interface{}{}})
 		} else {
 			allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
-			outputData, err := buildQueryResp(response.Results.TopK, response.Results.OutputFields, response.Results.FieldsData, response.Results.Ids, response.Results.Scores, allowJS)
+			outputData, err := buildQueryResp(searchResp.Results.TopK, searchResp.Results.OutputFields, searchResp.Results.FieldsData, searchResp.Results.Ids, searchResp.Results.Scores, allowJS)
 			if err != nil {
-				log.Warn("high level restful api, fail to deal with search result", zap.Any("result", response.Results), zap.Error(err))
+				log.Warn("high level restful api, fail to deal with search result", zap.Any("result", searchResp.Results), zap.Error(err))
 				c.JSON(http.StatusOK, gin.H{
 					HTTPReturnCode:    merr.Code(merr.ErrInvalidSearchResult),
 					HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),

internal/distributed/proxy/httpserver/handler_v1_test.go

@@ -2,6 +2,7 @@ package httpserver
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -14,6 +15,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"go.uber.org/atomic"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
@@ -21,6 +23,7 @@ import (
 	"github.com/milvus-io/milvus/internal/mocks"
 	"github.com/milvus-io/milvus/internal/proxy"
 	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
@@ -1789,3 +1792,43 @@ func getFieldSchema() []*schemapb.FieldSchema {
 
 	return fields
 }
+
+func TestInterceptor(t *testing.T) {
+	h := Handlers{}
+	v := atomic.NewInt32(0)
+	h.interceptors = []RestRequestInterceptor{
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			log.Info("pre1")
+			v.Add(1)
+			assert.EqualValues(t, 1, v.Load())
+			res, err := handler(ctx, req)
+			log.Info("post1")
+			v.Add(1)
+			assert.EqualValues(t, 6, v.Load())
+			return res, err
+		},
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			log.Info("pre2")
+			v.Add(1)
+			assert.EqualValues(t, 2, v.Load())
+			res, err := handler(ctx, req)
+			log.Info("post2")
+			v.Add(1)
+			assert.EqualValues(t, 5, v.Load())
+			return res, err
+		},
+		func(ctx context.Context, ginCtx *gin.Context, req any, handler func(reqCtx context.Context, req any) (any, error)) (any, error) {
+			log.Info("pre3")
+			v.Add(1)
+			assert.EqualValues(t, 3, v.Load())
+			res, err := handler(ctx, req)
+			log.Info("post3")
+			v.Add(1)
+			assert.EqualValues(t, 4, v.Load())
+			return res, err
+		},
+	}
+	_, _ = h.executeRestRequestInterceptor(context.Background(), nil, &milvuspb.CreateCollectionRequest{}, func(reqCtx context.Context, req any) (any, error) {
+		return &commonpb.Status{}, nil
+	})
+}

internal/distributed/proxy/service.go

@@ -192,7 +192,7 @@ func (s *Server) startHTTPServer(errChan chan error) {
 			return
 		}
 		c.Next()
-	}, authenticate, proxy.HTTPTraceLog)
+	}, authenticate)
 	app := ginHandler.Group("/v1")
 	httpserver.NewHandlers(s.proxy).RegisterRoutesToV1(app)
 	s.httpServer = &http.Server{Handler: ginHandler, ReadHeaderTimeout: time.Second}

internal/proxy/trace_log_interceptor.go

@@ -22,7 +22,6 @@ import (
 	"context"
 	"path"
 
-	"github.com/gin-gonic/gin"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 
@@ -44,6 +43,23 @@ func TraceLogInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryS
 		fields = append(fields, GetRequestFieldWithoutSensitiveInfo(req))
 		log.Ctx(ctx).Info("trace info: detail", fields...)
 		return handler(ctx, req)
+	case 3: // detail info with request and response
+		fields := GetRequestBaseInfo(ctx, req, info, true)
+		fields = append(fields, GetRequestFieldWithoutSensitiveInfo(req))
+		log.Ctx(ctx).Info("trace info: all request", fields...)
+		resp, err := handler(ctx, req)
+		if err != nil {
+			log.Ctx(ctx).Info("trace info: all, error", zap.Error(err))
+			return resp, err
+		}
+		if status, ok := requestutil.GetStatusFromResponse(resp); ok {
+			if status.Code != 0 {
+				log.Ctx(ctx).Info("trace info: all, fail", zap.Any("resp", resp))
+			}
+		} else {
+			log.Ctx(ctx).Info("trace info: all, unknown", zap.Any("resp", resp))
+		}
+		return resp, nil
 	default:
 		return handler(ctx, req)
 	}
@@ -94,8 +110,3 @@ func GetRequestFieldWithoutSensitiveInfo(req interface{}) zap.Field {
 	}
 	return zap.Any("request", req)
 }
-
-func HTTPTraceLog(ctx *gin.Context) {
-	// TODO trace http request info
-	ctx.Next()
-}

internal/proxy/trace_log_interceptor_test.go

@@ -24,9 +24,11 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
 
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
@@ -53,7 +55,7 @@ func TestTraceLogInterceptor(t *testing.T) {
 			DbName:         "db",
 			CollectionName: "col1",
 		}, &grpc.UnaryServerInfo{
-			FullMethod: "/milvus.proto.milvus.MilvusService/ShowCollections",
+			FullMethod: "/milvus.proto.milvus.MilvusService/CreateCollection",
 		}, handler)
 	}
 
@@ -64,7 +66,7 @@ func TestTraceLogInterceptor(t *testing.T) {
 			DbName:         "db",
 			CollectionName: "col1",
 		}, &grpc.UnaryServerInfo{
-			FullMethod: "/milvus.proto.milvus.MilvusService/ShowCollections",
+			FullMethod: "/milvus.proto.milvus.MilvusService/CreateCollection",
 		}, handler)
 	}
 
@@ -82,5 +84,51 @@ func TestTraceLogInterceptor(t *testing.T) {
 		})
 		assert.NotContains(t, strings.ToLower(fmt.Sprint(f2.Interface)), "password")
 	}
+
+	_ = paramtable.Get().Save(paramtable.Get().CommonCfg.TraceLogMode.Key, "3")
+	{
+		_, _ = TraceLogInterceptor(ctx, &milvuspb.CreateCollectionRequest{
+			DbName:         "db",
+			CollectionName: "col1",
+		}, &grpc.UnaryServerInfo{
+			FullMethod: "/milvus.proto.milvus.MilvusService/CreateCollection",
+		}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			return nil, errors.New("internet error")
+		})
+	}
+	{
+		_, _ = TraceLogInterceptor(ctx, &milvuspb.CreateCollectionRequest{
+			DbName:         "db",
+			CollectionName: "col1",
+		}, &grpc.UnaryServerInfo{
+			FullMethod: "/milvus.proto.milvus.MilvusService/CreateCollection",
+		}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			return &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError, Code: 500}, nil
+		})
+	}
+	{
+		_, _ = TraceLogInterceptor(ctx, &milvuspb.CreateCollectionRequest{
+			DbName:         "db",
+			CollectionName: "col1",
+		}, &grpc.UnaryServerInfo{
+			FullMethod: "/milvus.proto.milvus.MilvusService/CreateCollection",
+		}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			return "foo", nil
+		})
+	}
+	{
+		_, _ = TraceLogInterceptor(ctx, &milvuspb.ShowCollectionsRequest{
+			DbName: "db",
+		}, &grpc.UnaryServerInfo{
+			FullMethod: "/milvus.proto.milvus.MilvusService/ShowCollections",
+		}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			return &milvuspb.ShowCollectionsResponse{
+				Status: &commonpb.Status{
+					ErrorCode: commonpb.ErrorCode_Success,
+				},
+				CollectionNames: []string{"col1"},
+			}, nil
+		})
+	}
 	_ = paramtable.Get().Save(paramtable.Get().CommonCfg.TraceLogMode.Key, "0")
 }

internal/proxy/util.go

@@ -42,6 +42,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/commonpbutil"
+	"github.com/milvus-io/milvus/pkg/util/contextutil"
 	"github.com/milvus-io/milvus/pkg/util/crypto"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/metric"
@@ -899,12 +900,7 @@ func NewContextWithMetadata(ctx context.Context, username string, dbName string)
 	authKey := strings.ToLower(util.HeaderAuthorize)
 	authValue := crypto.Base64Encode(originValue)
 	dbKey := strings.ToLower(util.HeaderDBName)
-	contextMap := map[string]string{
-		authKey: authValue,
-		dbKey:   dbName,
-	}
-	md := metadata.New(contextMap)
-	return metadata.NewIncomingContext(ctx, md)
+	return contextutil.AppendToIncomingContext(ctx, authKey, authValue, dbKey, dbName)
 }
 
 func GetRole(username string) ([]string, error) {

pkg/util/contextutil/context_util.go

@@ -16,7 +16,12 @@
 
 package contextutil
 
-import "context"
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc/metadata"
+)
 
 type ctxTenantKey struct{}
 
@@ -37,3 +42,19 @@ func TenantID(ctx context.Context) string {
 
 	return ""
 }
+
+func AppendToIncomingContext(ctx context.Context, kv ...string) context.Context {
+	if len(kv)%2 == 1 {
+		panic(fmt.Sprintf("metadata: AppendToOutgoingContext got an odd number of input pairs for metadata: %d", len(kv)))
+	}
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.New(make(map[string]string, len(kv)/2))
+	}
+	for i, s := range kv {
+		if i%2 == 0 {
+			md.Append(s, kv[i+1])
+		}
+	}
+	return metadata.NewIncomingContext(ctx, md)
+}

pkg/util/contextutil/context_util_test.go

@@ -0,0 +1,44 @@
+/*
+ * 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 contextutil
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAppendToIncomingContext(t *testing.T) {
+	t.Run("invalid kvs", func(t *testing.T) {
+		assert.Panics(t, func() {
+			// nolint
+			AppendToIncomingContext(context.Background(), "foo")
+		})
+	})
+
+	t.Run("valid kvs", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = AppendToIncomingContext(ctx, "foo", "bar")
+		md, ok := metadata.FromIncomingContext(ctx)
+		assert.True(t, ok)
+		assert.Equal(t, "bar", md.Get("foo")[0])
+	})
+}

pkg/util/logutil/grpc_interceptor_test.go

@@ -61,5 +61,5 @@ func withMetaData(ctx context.Context, level zapcore.Level) context.Context {
 	md := metadata.New(map[string]string{
 		logLevelRPCMetaKey: level.String(),
 	})
-	return metadata.NewIncomingContext(context.TODO(), md)
+	return metadata.NewIncomingContext(ctx, md)
 }

pkg/util/requestutil/getter.go

@@ -140,6 +140,22 @@ func GetDSLFromRequest(req interface{}) (any, bool) {
 	return getter.GetDsl(), true
 }
 
+type StatusGetter interface {
+	GetStatus() *commonpb.Status
+}
+
+func GetStatusFromResponse(resp interface{}) (*commonpb.Status, bool) {
+	status, ok := resp.(*commonpb.Status)
+	if ok {
+		return status, true
+	}
+	getter, ok := resp.(StatusGetter)
+	if !ok {
+		return nil, false
+	}
+	return getter.GetStatus(), true
+}
+
 var TraceLogBaseInfoFuncMap = map[string]func(interface{}) (any, bool){
 	"collection_name": GetCollectionNameFromRequest,
 	"db_name":         GetDbNameFromRequest,

pkg/util/requestutil/getter_test.go

@@ -455,3 +455,56 @@ func TestGetDSLFromRequest(t *testing.T) {
 		})
 	}
 }
+
+func TestGetStatusFromResponse(t *testing.T) {
+	type args struct {
+		resp interface{}
+	}
+	tests := []struct {
+		name  string
+		args  args
+		want  *commonpb.Status
+		want1 bool
+	}{
+		{
+			name: "describe collection response",
+			args: args{
+				resp: &milvuspb.DescribeCollectionResponse{
+					Status: &commonpb.Status{
+						ErrorCode: commonpb.ErrorCode_Success,
+					},
+				},
+			},
+			want: &commonpb.Status{
+				ErrorCode: commonpb.ErrorCode_Success,
+			},
+			want1: true,
+		},
+		{
+			name: "common status",
+			args: args{
+				resp: &commonpb.Status{},
+			},
+			want:  &commonpb.Status{},
+			want1: true,
+		},
+		{
+			name: "invalid response",
+			args: args{
+				resp: "foo",
+			},
+			want1: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, got1 := GetStatusFromResponse(tt.args.resp)
+			if got1 && !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GetStatusFromResponse() got = %v, want %v", got, tt.want)
+			}
+			if got1 != tt.want1 {
+				t.Errorf("GetStatusFromResponse() got1 = %v, want %v", got1, tt.want1)
+			}
+		})
+	}
+}