responsex/response.go

package responsex

// Author: Yun
// Version: 2023年6月12日19:25:54

import (
	"context"
	"encoding/json"
	"math"
	"net/http"

	"github.com/yuninks/errorx"
	"github.com/yuninks/langx"
	"github.com/zeromicro/go-zero/core/trace"
)

type response struct {
	Code       int         `json:"code"`
	Message    string      `json:"message"`
	Data       any         `json:"data,omitempty"`
	Pagination *Pagination `json:"pagination,omitempty"`
	TraceId    string      `json:"trace_id,omitempty"`
}

type Pagination struct {
	Page       int64 `json:"page"`
	Size       int64 `json:"size"`
	TotalPage  int64 `json:"total_page"`
	TotalCount int64 `json:"total_count"`
}

// 计算DB分页偏移量
func GetOffset(page, size int) (offset int, limit int) {
	if page <= 0 {
		page = 1
	}
	offset = (page - 1) * size
	limit = size
	if limit == 0 {
		// 默认查10条
		limit = 10
	}
	return
}

// 计算响应分页
func GetPage(page, size, totalCount int64) Pagination {
	pa := Pagination{}
	totalPage := 0
	if size != 0 {
		totalPage = int(math.Ceil(float64(totalCount) / float64(size)))
	}
	pa.Page = page
	pa.Size = size
	pa.TotalPage = int64(totalPage)
	pa.TotalCount = totalCount
	return pa
}

// Msg格式化响应
func FormatMessage(ctx context.Context, w http.ResponseWriter, msgKey string, format map[string]string, data interface{}, ops ...Option) {
	code, msg := langx.GetTransFormatCtx(ctx, msgKey, format)
	baseOps := []Option{SetCode(code), SetMessage(msg), SetData(data)}
	baseOps = append(baseOps, ops...)
	ResponseCtx(ctx, w, baseOps...)
}

// 成功的响应
func Success(ctx context.Context, w http.ResponseWriter, data any, ops ...Option) {
	baseOps := []Option{SetCode(op.defaultSuccessCode), SetMessage("请求成功"), SetData(data)}
	baseOps = append(baseOps, ops...)
	ResponseCtx(ctx, w, baseOps...)
}

func SuccessWithPage(ctx context.Context, w http.ResponseWriter, data any, page *Pagination, ops ...Option) {
	baseOps := []Option{SetCode(op.defaultSuccessCode), SetMessage("请求成功"), SetData(data), SetPage(page)}
	baseOps = append(baseOps, ops...)
	ResponseCtx(ctx, w, baseOps...)
}

// 失败的响应
func Error(ctx context.Context, w http.ResponseWriter, err error, ops ...Option) {
	ErrorWithData(ctx, w, err, "", ops...)
}

func ErrorWithData(ctx context.Context, w http.ResponseWriter, err error, data interface{}, ops ...Option) {
	code := op.defaultErrorCode
	msg := "请求失败"
	if err != nil {
		val,ok := errorx.As(err)
		if ok {
			
			if langx.GetDefaultCode() != val.Code() {
				code = val.Code()
			}
		}
		msg = err.Error()
	}
	baseOps := []Option{SetCode(code), SetMessage(msg), SetData(data)}
	baseOps = append(baseOps, ops...)
	ResponseCtx(ctx, w, baseOps...)
}

func ErrorStr(ctx context.Context, w http.ResponseWriter, msg string, ops ...Option) {
	msg = langx.GetMsgCtx(ctx, msg)
	code := langx.GetCode(msg)
	if code == langx.GetDefaultCode() {
		code = op.defaultErrorCode
	}
	baseOps := []Option{SetCode(code), SetMessage(msg)}
	baseOps = append(baseOps, ops...)
	ResponseCtx(ctx, w, baseOps...)
}

// 基础的响应
func ResponseCtx(ctx context.Context, w http.ResponseWriter, ops ...Option) {

	def := defaultOptions() // 默认值
	for _, o := range ops {
		o(def)
	}

	resp := response{}
	resp.Code = def.code
	resp.Message = def.message
	resp.Data = def.data
	resp.Pagination = def.pagination

	// 获取trace_id
	if def.traceId != "" {
		t := ctx.Value(def.traceId)
		if t != nil {
			val, ok := t.(string)
			if ok {
				resp.TraceId = val
			}
		} else {
			// go-zero
			resp.TraceId = trace.TraceIDFromContext(ctx)
		}
	}

	bs, _ := json.Marshal(resp)

	if !def.ignoreLog {
		// 记录日志
		def.logger.Info(ctx, "response:", string(bs))
	}

	// 写入自定义响应头
	for k, v := range def.headers {
		w.Header().Set(k, v)
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	// 确定HTTP状态码
	httpStatus := def.httpStatusCode
	if httpStatus == 0 {
		httpStatus = 200
	}
	w.WriteHeader(httpStatus)

	n, err := w.Write(bs)
	if err != nil || n != len(bs) {
		op.logger.Error(ctx, "write response failed err:%+v", err)
	}
}