bbgo/pkg/notifier/larknotifier/lark.go
2024-08-30 23:03:08 +08:00

197 lines
3.9 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package larknotifier
import (
"bytes"
"context"
"encoding/json"
"fmt"
"git.qtrade.icu/lychiyu/bbgo/pkg/types"
"github.com/sirupsen/logrus"
"golang.org/x/time/rate"
"net/http"
"reflect"
"time"
)
var apiLimiter = rate.NewLimiter(rate.Every(time.Second), 5)
var log = logrus.WithField("service", "lark")
type notifyTask struct {
message string
}
type Notifier struct {
token string
taskC chan notifyTask
}
type Option func(notifier *Notifier)
type TextMessage struct {
MsgType string `json:"msg_type"`
Content struct {
Text string `json:"text"`
} `json:"content"`
}
func getLocDateTimeString() string {
loc, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
fmt.Println("Error loading location:", err)
return ""
}
// 获取当前东八区时间
currentTime := time.Now().In(loc)
// 格式化时间
formattedTime := currentTime.Format("2006-01-02 15:04:05")
return formattedTime
}
func sendLarkMessage(token, text string) error {
// 创建消息结构体
message := TextMessage{
MsgType: "text",
Content: struct {
Text string `json:"text"`
}{
Text: fmt.Sprintf("通知时间:%s\n内容\n%s", getLocDateTimeString(), text),
},
}
// 将消息结构体序列化为 JSON
jsonData, err := json.Marshal(message)
if err != nil {
return fmt.Errorf("failed to marshal JSON: %v", err)
}
webhookURL := fmt.Sprintf("https://open.larksuite.com/open-apis/bot/v2/hook/%s", token)
// 创建 HTTP POST 请求
req, err := http.NewRequest("POST", webhookURL, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %v", err)
}
req.Header.Set("Content-Type", "application/json")
// 发送 HTTP 请求
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %v", err)
}
defer resp.Body.Close()
// 检查响应状态码
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("request failed with status code: %d", resp.StatusCode)
}
return nil
}
func New(token string, options ...Option) *Notifier {
notifier := &Notifier{
token: token,
taskC: make(chan notifyTask, 100),
}
for _, o := range options {
o(notifier)
}
go notifier.worker()
return notifier
}
func (n *Notifier) worker() {
ctx := context.Background()
for {
select {
case <-ctx.Done():
return
case task := <-n.taskC:
apiLimiter.Wait(ctx)
n.consume(task)
}
}
}
func (n *Notifier) consume(task notifyTask) {
if task.message != "" {
if err := sendLarkMessage(n.token, task.message); err != nil {
log.WithError(err).Error("lark send error")
}
}
}
func (n *Notifier) Notify(obj interface{}, args ...interface{}) {
n.NotifyTo("", obj, args...)
}
func filterPlaintextMessages(args []interface{}) (texts []string, pureArgs []interface{}) {
var firstObjectOffset = -1
for idx, arg := range args {
rt := reflect.TypeOf(arg)
if rt.Kind() == reflect.Ptr {
switch a := arg.(type) {
case nil:
texts = append(texts, "nil")
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
case types.PlainText:
texts = append(texts, a.PlainText())
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
case types.Stringer:
texts = append(texts, a.String())
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
}
}
}
pureArgs = args
if firstObjectOffset > -1 {
pureArgs = args[:firstObjectOffset]
}
return texts, pureArgs
}
func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}) {
var message string
switch a := obj.(type) {
case string:
message = fmt.Sprintf(a, args...)
default:
//log.Errorf("unsupported notification format: %T %+v", a, a)
return
}
select {
case n.taskC <- notifyTask{
message: message,
}:
default:
log.Error("[lark] cannot send task to notify")
}
}
func (n *Notifier) SendPhoto(buffer *bytes.Buffer) {
n.SendPhotoTo("", buffer)
}
func (n *Notifier) SendPhotoTo(channel string, buffer *bytes.Buffer) {
// TODO
}