wss-pool/cmd/websocketcollect/cache/leveldb.go

package cache

import (
	"encoding/json"
	"fmt"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/filter"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"go.mongodb.org/mongo-driver/bson"
	"strings"
	"wss-pool/config"
	"wss-pool/internal"
	"wss-pool/internal/data"
	"wss-pool/logging/applogger"
	"wss-pool/pkg/model"
	"wss-pool/pkg/model/stock"
)

var db *leveldb.DB

// initDB
//
//	@Description:
func InitDB(name string) {
	o := &opt.Options{
		Filter: filter.NewBloomFilter(10),
	}

	var tableNamePath string
	switch config.Config.ServerLevel {
	case "test":
		tableNamePath = fmt.Sprintf("./cmd/websocketcollect/db/%v", name)
	default:
		tableNamePath = fmt.Sprintf("/home/ubuntu/wss-server/db/%v", name)
	}

	dba, err := leveldb.OpenFile(tableNamePath, o)
	if err != nil {
		applogger.Error("OpenFile err:%v", err)
		return
	}

	db = dba
}

// WriteDB
//
//	@Description:
//	@param key
//	@param value
//	@return error
func WriteDB(key, value string) error {
	batch := new(leveldb.Batch)
	batch.Put([]byte(key), []byte(value))
	if err := db.Write(batch, nil); err != nil {
		fmt.Printf("Write err:%v\n", err)
		return err
	}

	return nil
}

// ReadDB
//
//	@Description:
//	@param key
//	@return []byte
func ReadDB(key string) []byte {
	data, err := db.Get([]byte(key), nil)
	if err != nil {
		fmt.Printf("OpenFile err:%v\n", err)
		return []byte{}
	}

	return data
}

// ReadListDB
//
//	@Description:
//	@return []string
func ReadListDB() []string {
	var privateKey []string
	iter := db.NewIterator(nil, nil)
	for iter.Next() {
		value := iter.Value()
		privateKey = append(privateKey, string(value))
	}
	iter.Release()
	if err := iter.Error(); err != nil {
		fmt.Printf("iter.Error:%v\n", err)
		return privateKey
	}

	return privateKey
}

// DeleteByKey
//
//	@Description:
//	@param key
//	@return error
func DeleteByKey(key string) error {
	if err := db.Delete([]byte(key), nil); err != nil {
		return err
	}

	return nil
}

// InitShareUsCode
//
//	@Description: 初始化美股股票代码列表
func InitShareUsCode() {
	filter := bson.M{"Country": "US", "YesterdayClose": bson.M{"$ne": ""}}
	projection := bson.M{"Code": 1}

	for _, mongoClient := range data.MgoDbClientMap {
		stockRes := make([]stock.StockPolygon, 0)
		data.MgoPagingFindStructProjectionNew(mongoClient, data.StockList, filter, projection, 20000, 1, -1, &stockRes)
		for _, v := range stockRes {
			if err := WriteDB(v.Code, v.Code); err != nil {
				applogger.Error("sendUsCodeNew err:%v", err)
				continue
			}
		}
		applogger.Info("写入美股股票代码列表:%v", ReadListDB())
		applogger.Debug("写入美股股票代码列表总数:%v", len(ReadListDB()))
	}
}

// InitForexCode
//
//	@Description: 初始化外汇代码列表
func InitForexCode() {
	// https://finnhub.io/api/v1/forex/symbol?exchange=oanda&token=cqt0409r01qvdch2gnpgcqt0409r01qvdch2gnq0
	url := fmt.Sprintf("https://%vforex/symbol?exchange=oanda&token=%v", config.Config.FinnhubUs.FinnhubHost, config.Config.FinnhubUs.FinnhubKey)
	applogger.Debug("url data info：%v", url)

	bodyStr, err := internal.HttpGet(url)
	if err != nil {
		applogger.Error("WriteShareUs err info:%v", err)
		return
	}
	applogger.Debug("new add shareUs code info：%v", bodyStr)

	var dataForex []model.ForexCodeList
	if err = json.Unmarshal([]byte(bodyStr), &dataForex); err != nil {
		applogger.Error("json.Unmarshal err info:%v", err)
		return
	}
	for _, v := range dataForex {
		code := strings.Replace(v.DisplaySymbol, "/", "", -1)
		if err := WriteDB(code, code); err != nil {
			applogger.Error("sendUsCodeNew err:%v", err)
			continue
		}
	}
	applogger.Info("写入外汇代码列表:%v", ReadListDB())
	applogger.Debug("写入外汇代码列表总数:%v", len(ReadListDB()))
}