esproxy/internal/service/middle.go

package service

import (
	"bytes"
	"context"
	"esproxy/internal/consts"
	"fmt"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/util/gconv"
	"io"
	"time"
)

// Middleware 分析查询记录日志中间件
// 获取请求时间、查询和响应
// 记录查询到数据库、区分查询类型：简单查询、聚合查询、复杂查询
// 根据es当前状态，使用4个通道，进行调度
func Middleware(r *ghttp.Request) {
	now := time.Now()
	bodyBytes, err := io.ReadAll(r.Request.Body)
	if len(bodyBytes) > 0 && err == nil {
		finalBytes := bodyBytes
		r.ContentLength = gconv.Int64(len(finalBytes))
		r.Request.Header.Set("Content-Length", fmt.Sprintf("%d", len(finalBytes)))
		r.Request.Body = io.NopCloser(bytes.NewReader(finalBytes))
	}
	queryLevel := getQueryLevel(bodyBytes)
	r.SetCtxVar(consts.QueryLevelKey, queryLevel)
	r.Middleware.Next()
	if r.RequestURI != "/" { //打印非建立链接的日志
		g.Log().Infof(r.Context(), "status:%d  time:%fs  req:%s   waitPool:%fs  level:%d  from:%s  query:%s", r.Response.Status, time.Since(now).Seconds(), r.RequestURI, r.GetCtxVar(consts.QueryWaitPoolTime).Float64(), queryLevel, r.GetClientIp(), string(bodyBytes))
	}
}

var (
	aggsFlag = []byte("aggs")
)

// getQueryLevel 获取sql查询复杂度
func getQueryLevel(detail []byte) consts.QueryLevel {
	if isAggs := bytes.Contains(detail, aggsFlag); !isAggs {
		return consts.QueryLevelSimple
	}
	if queryLen := len(string(detail)); queryLen < g.Cfg().MustGet(context.Background(), "elasticSearch.complexQueryLen", 200).Int() {
		return consts.QueryLevelAggs
	}
	return consts.QueryLevelComplex
}