sparrow/pkg/ruleEngine/nodes/rest_api_request_node.go

package nodes

import (
	"bytes"
	"encoding/json"
	"github.com/gogf/gf/os/grpool"
	"net/http"
	"sparrow/pkg/protocol"
	"sparrow/pkg/ruleEngine"
	"sparrow/pkg/server"
	"sparrow/pkg/utils"
	"sync"
	"time"
)

// RestApiRequestNode 请求外部API节点
// 支持的配置字符
/*
{
    "url": "http://localhost/api/test",
    "headers": {},
    "retry": 3,
    "method": "post",
    "time_out": 5,
    "retry_wait": 1
}
*/
type RestApiRequestNode struct {
	pool   *grpool.Pool
	config *RestApiRequestNodeConfig
	client *utils.HttpClient
	bufPool sync.Pool
}

func (r *RestApiRequestNode) Init(ctx ruleEngine.Context, config string) error {
	r.pool = grpool.New(500)
	if config == "" {
		r.config = &RestApiRequestNodeConfig{
			Url:       "http://localhost:8899/test",
			Headers:   make(map[string]string),
			Retry:     1,
			Method:    "POST",
			TimeOut:   5,
			RetryWait: 1,
		}
	} else {
		c := new(RestApiRequestNodeConfig)
		err := json.Unmarshal([]byte(config), c)
		if err != nil {
			return err
		}
		r.config = c
	}
	client := utils.NewHttpClientWithConfig(time.Duration(r.config.TimeOut)*time.Second,
		r.config.Retry, time.Duration(r.config.RetryWait)*time.Second)
	client.SetLogger(server.Log)
	r.client = client
	r.bufPool = sync.Pool{
		New: func() interface{}{
			return new(bytes.Buffer)
		},
	}
	return nil
}

func (r *RestApiRequestNode) OnMessage(ctx ruleEngine.Context, message *protocol.Message) error {
	body := message.Data
	headers := make(map[string]string)
	for k, v := range message.MetaData {
		headers[k] = v.(string)
	}
	for k, v := range r.config.Headers {
		headers[k] = v
	}
	w := r.bufPool.Get().(*bytes.Buffer)
	if err := json.NewEncoder(w).Encode(body); err != nil {
		return err
	}
	w.Reset()
	r.bufPool.Put(w)
	req, err := utils.NewRequest(r.config.Method, r.config.Url, []byte(body))
	if err != nil {
		server.Log.Error(err)
		return err
	}
	for k, v := range headers {
		req.Header.Add(k, v)
	}
	req.Header.Add("Content-Type", "application/json")
	return r.pool.Add(func() {
		res, err := r.client.Do(req)
		if err != nil {
			server.Log.Errorf("请求出错%s", err.Error())
			next := r.processError(ctx, message, err)
			ctx.TellError(next, err)
			return
		}
		if res != nil && res.Response() != nil {
			if res.Response().StatusCode == http.StatusOK {
				msg := r.processResponse(ctx, message, res)
				ctx.TellSuccess(msg)
			}
		}
	})
}

func (r *RestApiRequestNode) processError(ctx ruleEngine.Context, msg *protocol.Message, err error) *protocol.Message {
	var metaData map[string]interface{}
	if msg.MetaData != nil {
		metaData = msg.MetaData
	}
	metaData = make(map[string]interface{})
	metaData["error"] = err.Error()

	return ctx.TransformMessage(msg, msg.Type, msg.Originator, metaData, msg.Data)
}
func (r *RestApiRequestNode) processResponse(ctx ruleEngine.Context, msg *protocol.Message, res utils.Responser) *protocol.Message {
	var metaData map[string]interface{}
	if msg.MetaData != nil {
		metaData = msg.MetaData
	}
	metaData = make(map[string]interface{})
	metaData["STATUS"] = res.Response().Status
	metaData["STATUS_CODE"] = res.Response().StatusCode
	var data string
	if body, err := res.String(); err == nil {
		data = body
	}
	return ctx.TransformMessage(msg, msg.Type, msg.Originator, metaData, data)
}

type RestApiRequestNodeConfig struct {
	Url       string            `json:"url"`        // 请求的地址
	Headers   map[string]string `json:"headers"`    // 请求头
	Retry     int               `json:"retry"`      // 重试次数
	Method    string            `json:"method"`     // 请求方法
	TimeOut   int               `json:"time_out"`   // 超时时间(秒)
	RetryWait int               `json:"retry_wait"` // 重试等待时间
}