wss-pool/internal/data/business/mgohistoricalus.go

package business

import (
	"encoding/json"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"strconv"
	"strings"
	"sync"
	"wss-pool/config"
	"wss-pool/dictionary"
	"wss-pool/internal"
	"wss-pool/internal/data"
	"wss-pool/logging/applogger"
	"wss-pool/pkg/model/stock"
)

/*
采集美股历史数据
*/
const (
	UsPageSize int64 = 500
)

func UpdateStockUs(start, end interface{}, state string) {
	fmt.Println("updateStockUs", start, end, state)
	stocks, _, pageTotal := GetStockAll("US", 1, UsPageSize)
	wg := sync.WaitGroup{}
	for _, value := range stocks {
		wg.Add(1)
		go func(start, end interface{}, state, code string) {
			defer wg.Done()
			param := fmt.Sprintf("apiKey=%v", config.Config.ShareGather.PolygonKey)
			multiplier, _ := strconv.Atoi(dictionary.StockUsListTimeMap[state])
			timespan := "minute"
			if !strings.Contains(state, "min") {
				multiplier = 1
				timespan = dictionary.StockUsListTimeMap[state]
			}
			url := fmt.Sprintf("https://%v/v2/aggs/ticker/%v/range/%d/%v/%v/%v?adjusted=true&sort=asc&%v",
				config.Config.ShareGather.PolygonHost, code, multiplier, timespan, start, end, param)
			bodyStr, err := internal.HttpGet(url)
			if err != nil {
				applogger.Error("Failed to query data：%v", err)
				return
			}
			var eodModel stock.PreviousCloseResponse
			if err = json.Unmarshal([]byte(bodyStr), &eodModel); err != nil {
				applogger.Error("eodModel json Unmarshal err: %v", err)
				return
			}
			applogger.Info("info", eodModel)
			UpdateUs(eodModel, state)
		}(start, end, state, value.Code)
	}
	wg.Wait()
	for i := int64(2); i <= pageTotal; i++ {
		stocks, _, _ := GetStockAll("US", i, UsPageSize)
		wg := sync.WaitGroup{}
		for _, value := range stocks {
			wg.Add(1)
			go func(start, end interface{}, state, code string) {
				defer wg.Done()
				param := fmt.Sprintf("apiKey=%v", config.Config.ShareGather.PolygonKey)
				multiplier, _ := strconv.Atoi(dictionary.StockUsListTimeMap[state])
				timespan := "minute"
				if !strings.Contains(state, "min") {
					multiplier = 1
					timespan = dictionary.StockUsListTimeMap[state]
				}
				url := fmt.Sprintf("https://%v/v2/aggs/ticker/%v/range/%d/%v/%v/%v?adjusted=true&sort=asc&%v",
					config.Config.ShareGather.PolygonHost, code, multiplier, timespan, start, end, param)
				bodyStr, err := internal.HttpGet(url)
				if err != nil {
					applogger.Error("Failed to query data：%v", err)
					return
				}
				var eodModel stock.PreviousCloseResponse
				if err = json.Unmarshal([]byte(bodyStr), &eodModel); err != nil {
					applogger.Error("eodModel json Unmarshal err: %v", err)
					return
				}
				UpdateUs(eodModel, state)
				applogger.Info("info", eodModel)
			}(start, end, state, value.Code)
		}
		wg.Wait()
	}
}

func UpdateUs(result stock.PreviousCloseResponse, period string) {
	if len(result.Results) <= 0 {
		applogger.Error("us data is null %v", result)
		return
	}
	var dataList []mongo.WriteModel
	for _, val := range result.Results {
		filter := bson.M{"code": bson.M{"$eq": result.Ticker}, "timestamp": bson.M{"$eq": val.T}}
		update := bson.D{{"$set", bson.D{
			{"code", result.Ticker},
			{"timestamp", val.T},
			{"vw", val.VW.String()},
			{"o", val.O.String()},
			{"c", val.C.String()},
			{"h", val.H.String()},
			{"n", val.N},
		}}}
		applogger.Info("UpdateUs info:%v", update)
		models := mongo.NewUpdateOneModel().SetFilter(filter).SetUpdate(update).SetUpsert(true)
		dataList = append(dataList, models)
	}
	tableName := data.GetStockUsTableName(period)
	if len(dataList) > 0 {
		if err := data.MgoBulkWrite(tableName, dataList); err != nil {
			applogger.Error("UpdateUs MgoInsertMany err:%v", err)
		}
	}
}