baize/framework/core/api/context.go

package api

import (
	"bytes"
	"encoding/json"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/logger"
	"github.com/gin-gonic/gin"
	"github.com/pkg/errors"
	"io"
	"mime/multipart"
	"strings"
)

const (
	bodyKey        = "body-context"
	queryParamsKey = "query-params-context"
	pathParamsKey  = "path-params-context"
)

type Context struct {
	*gin.Context
}

// GetFileHeaderBytes 获取传递的文件名和文件内容
func (c *Context) GetFileHeaderBytes(fileHeader *multipart.FileHeader) (string, []byte, error) {
	file, err := fileHeader.Open()
	if err != nil {
		return "", nil, errors.New(err.Error())
	}

	defer func(file multipart.File) {
		err := file.Close()
		if err != nil {
			logger.GetInstance().Error(errors.New(err.Error()))
			return
		}
	}(file)

	contentBytes, err := io.ReadAll(file)
	if err != nil {
		return "", nil, errors.New(err.Error())
	}

	return fileHeader.Filename, contentBytes, nil
}

func (c *Context) GetHeaders() map[string]string {
	headers := make(map[string]string)

	for key, values := range c.Request.Header {
		headers[key] = strings.Join(values, ",")
	}

	return headers
}

type CacheBody struct {
	c         *Context
	bytesBody []byte
}

func (cacheBody *CacheBody) Set(bytesBody []byte) {
	cacheBody.bytesBody = bytesBody
	cacheBody.c.Set(bodyKey, cacheBody.bytesBody)
}

func (cacheBody *CacheBody) Bytes() []byte {
	return cacheBody.bytesBody
}

func (cacheBody *CacheBody) Marshal(output any) error {
	bytesBody, err := json.Marshal(output)
	if err != nil {
		return errors.New(err.Error())
	}

	cacheBody.Set(bytesBody)

	return nil
}

func (cacheBody *CacheBody) Unmarshal(output any) error {
	err := json.Unmarshal(cacheBody.Bytes(), output)
	if err != nil {
		return errors.New(err.Error())
	}

	return nil
}

func (c *Context) GetBytesBody() (*CacheBody, error) {
	body, exist := c.Get(bodyKey)
	if !exist {
		bytesBody, err := c.readOriginBody()
		if err != nil {
			return nil, err
		}

		return &CacheBody{
			c:         c,
			bytesBody: bytesBody,
		}, nil
	}

	switch b := body.(type) {
	case []byte:
		return &CacheBody{
			c:         c,
			bytesBody: b,
		}, nil
	case map[string]any:
		bytesBody, err := json.Marshal(b)
		if err != nil {
			return nil, errors.New(err.Error())
		}

		return &CacheBody{
			c:         c,
			bytesBody: bytesBody,
		}, nil
	default:
		return nil, errors.New("不支持的body类型")
	}
}

type JsonBody struct {
	c           *Context
	jsonBodyMap map[string]any
}

func (jsonBody *JsonBody) Set(key string, value any) {
	jsonBody.jsonBodyMap[key] = value
	jsonBody.c.Set(bodyKey, jsonBody.jsonBodyMap)
}

func (jsonBody *JsonBody) Delete(key string) {
	delete(jsonBody.jsonBodyMap, key)
	jsonBody.c.Set(bodyKey, jsonBody.jsonBodyMap)
}

func (jsonBody *JsonBody) Get(key string) any {
	return jsonBody.jsonBodyMap[key]
}

func (jsonBody *JsonBody) Map() map[string]any {
	return jsonBody.jsonBodyMap
}

func (jsonBody *JsonBody) Bytes() ([]byte, error) {
	jsonBytes, err := json.Marshal(jsonBody.jsonBodyMap)
	if err != nil {
		return nil, errors.New(err.Error())
	}

	return jsonBytes, nil
}

func (jsonBody *JsonBody) Unmarshal(output any) error {
	jsonBytes, err := jsonBody.Bytes()
	if err != nil {
		return err
	}

	err = json.Unmarshal(jsonBytes, output)
	if err != nil {
		return errors.New(err.Error())
	}

	return nil
}

func (c *Context) GetJsonBody() (*JsonBody, error) {
	body, exist := c.Get(bodyKey)
	if !exist {
		bytesBody, err := c.readOriginBody()
		if err != nil {
			return nil, err
		}

		jsonBodyMap := make(map[string]any)
		err = json.Unmarshal(bytesBody, &jsonBodyMap)
		if err != nil {
			return nil, errors.New(err.Error())
		}

		return &JsonBody{
			c:           c,
			jsonBodyMap: jsonBodyMap,
		}, nil
	}

	switch b := body.(type) {
	case []byte:
		jsonBodyMap := make(map[string]any)
		err := json.Unmarshal(b, &jsonBodyMap)
		if err != nil {
			return nil, errors.New(err.Error())
		}

		return &JsonBody{
			c:           c,
			jsonBodyMap: jsonBodyMap,
		}, nil
	case map[string]any:
		return &JsonBody{
			c:           c,
			jsonBodyMap: body.(map[string]any),
		}, nil
	default:
		return nil, errors.New("不支持的body类型")
	}
}

type QueryPrams struct {
	c           *Context
	queryParams map[string]string
}

func (queryParams *QueryPrams) Set(key string, value string) {
	queryParams.queryParams[key] = value
	queryParams.c.Set(queryParamsKey, queryParams.queryParams)
}

func (queryParams *QueryPrams) Delete(key string) {
	delete(queryParams.queryParams, key)
	queryParams.c.Set(queryParamsKey, queryParams.queryParams)
}

func (queryParams *QueryPrams) Get(key string) string {
	return queryParams.queryParams[key]
}

func (queryParams *QueryPrams) Map() map[string]string {
	return queryParams.queryParams
}

func (c *Context) GetQueryParams() *QueryPrams {
	queryParams, exist := c.Get(queryParamsKey)
	if !exist {
		return &QueryPrams{
			c:           c,
			queryParams: c.getAllQueryParams(),
		}
	}

	return &QueryPrams{
		c:           c,
		queryParams: queryParams.(map[string]string),
	}
}

type PathPrams struct {
	c          *Context
	pathParams map[string]string
}

func (pathParams *PathPrams) Set(key string, value string) {
	pathParams.pathParams[key] = value
	pathParams.c.Set(pathParamsKey, pathParams.pathParams)
}

func (pathParams *PathPrams) Delete(key string) {
	delete(pathParams.pathParams, key)
	pathParams.c.Set(pathParamsKey, pathParams.pathParams)
}

func (pathParams *PathPrams) Get(key string) string {
	return pathParams.pathParams[key]
}

func (pathParams *PathPrams) Map() map[string]string {
	return pathParams.pathParams
}

func (c *Context) GetPathParams() *PathPrams {
	pathParams, exist := c.Get(pathParamsKey)
	if !exist {
		return &PathPrams{
			c:          c,
			pathParams: c.getAllPathParams(),
		}
	}

	return &PathPrams{
		c:          c,
		pathParams: pathParams.(map[string]string),
	}
}

func (c *Context) getAllQueryParams() map[string]string {
	queryParams := make(map[string]string)

	for key, values := range c.Request.URL.Query() {
		queryParams[key] = strings.Join(values, ",")
	}

	return queryParams
}

func (c *Context) getAllPathParams() map[string]string {
	pathParams := make(map[string]string)

	for _, params := range c.Params {
		pathParams[params.Key] = params.Value
	}

	return pathParams
}

func (c *Context) readOriginBody() ([]byte, error) {
	if c.Request.Body == nil {
		return make([]byte, 0), nil
	}

	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return nil, errors.New(err.Error())
	}

	defer func(Body io.ReadCloser) {
		err := Body.Close()
		if err != nil {
			logger.GetInstance().Error(errors.New(err.Error()))
			return
		}
	}(c.Request.Body)

	c.Request.Body = io.NopCloser(bytes.NewBuffer(body))

	c.Set(bodyKey, body)

	return body, nil
}

const (
	tenantInfoKey = "context-tenant-info"
	userInfoKey   = "context-user-info"
)

type TenantInfo interface {
	GetID() string
	GetName() string
}

type UserInfo interface {
	GetID() string
	GetUserName() string
	GetName() string
}

func (c *Context) SetTenantInfo(tenantInfo TenantInfo) {
	c.Set(tenantInfoKey, tenantInfo)
}

func (c *Context) SetUserInfo(userInfo UserInfo) {
	c.Set(userInfoKey, userInfo)
}

func (c *Context) GetTenantInfo() TenantInfo {
	tenantInfo, exist := c.Get(tenantInfoKey)
	if !exist {
		return nil
	}

	return tenantInfo.(TenantInfo)
}

func (c *Context) GetUserInfo() UserInfo {
	userInfo, exist := c.Get(userInfoKey)
	if !exist {
		return nil
	}

	return userInfo.(UserInfo)
}