eta_crawler/models/base_from_eia_steo.go

package models

import (
	"eta/eta_crawler/utils"
	"time"

	"github.com/beego/beego/v2/client/orm"
	"github.com/shopspring/decimal"
)

// BaseFromEiaSteoIndex EiaSteo指标
type BaseFromEiaSteoIndex struct {
	BaseFromEiaSteoIndexId    int       `orm:"column(base_from_eia_steo_index_id);pk"`
	BaseFromEiaSteoClassifyId int       `description:"指标分类id"`
	IndexCode                 string    `description:"指标编码"`
	IndexName                 string    `description:"指标名称（中文名称）"`
	IndexNameOriginal         string    `description:"指标名称（原始名称）"`
	Frequency                 string    `description:"频度"`
	Level                     int       `description:"指标层级（原始层级）"`
	Unit                      string    `description:"单位"`
	Super                     string    `description:"我也不知道是个啥，反正先存起来，万一用到了呢"`
	Precision                 int       `description:"精度，我也不知道会不会用到，反正先存起来，万一用到了呢"`
	LastHistorical            string    `description:"最近的历史记录，我也不知道会不会用到，反正先存起来，万一用到了呢"`
	Description               string    `description:"备注信息，我也不知道会不会用到，反正先存起来，万一用到了呢"`
	IsMappable                int       `description:"是否可映射，我也不知道会不会用到，反正先存起来，万一用到了呢"`
	StartDate                 time.Time `description:"开始日期"`
	EndDate                   time.Time `description:"结束日期"`
	ModifyTime                time.Time `description:"最新更新时间"`
	CreateTime                time.Time `description:"创建时间"`
}

// Add 新增指标
func (item *BaseFromEiaSteoIndex) Add() (err error) {
	o := orm.NewOrmUsingDB("data")
	lastId, err := o.Insert(item)
	if err != nil {
		return
	}
	item.BaseFromEiaSteoIndexId = int(lastId)
	return
}

// Update 更新作者基础信息
func (item *BaseFromEiaSteoIndex) Update(cols []string) (err error) {
	o := orm.NewOrmUsingDB("data")
	_, err = o.Update(item, cols...)
	return
}

// BaseFromEiaSteoClassify EiaSteo分类
type BaseFromEiaSteoClassify struct {
	BaseFromEiaSteoClassifyId int       `orm:"column(base_from_eia_steo_classify_id);pk"`
	ClassifyName              string    `description:"分类名称（中文名称）"`
	ClassifyNameOriginal      string    `description:"分类名称（原始名称）"`
	ParentId                  int       `description:"父级id"`
	Level                     int       `description:"层级"`
	ModifyTime                time.Time `description:"最新更新时间"`
	CreateTime                time.Time `description:"创建时间"`
}

// 新增分类
func (item *BaseFromEiaSteoClassify) AddBaseFromEiaSteoClassify() (err error) {
	o := orm.NewOrmUsingDB("data")
	lastId, err := o.Insert(item)
	if err != nil {
		return
	}
	item.BaseFromEiaSteoClassifyId = int(lastId)
	return
}

// BaseFromEiaSteoData EiaSteo数据表
type BaseFromEiaSteoData struct {
	BaseFromEiaSteoDataId  int       `orm:"column(base_from_eia_steo_data_id);pk"`
	BaseFromEiaSteoIndexId int       `description:"指标id"`
	IndexCode              string    `description:"指标编码"`
	DataTime               time.Time `description:"数据日期"`
	Value                  float64   `description:"数据值"`
	ModifyTime             time.Time `description:"最新更新时间"`
	CreateTime             time.Time `description:"创建时间"`
}

// GetBaseFromEiaSteoClassifyAll 获取所有的分类
func GetBaseFromEiaSteoClassifyAll() (list []*BaseFromEiaSteoClassify, err error) {
	o := orm.NewOrmUsingDB("data")
	sql := `SELECT * FROM base_from_eia_steo_classify`
	_, err = o.Raw(sql).QueryRows(&list)
	return
}

func GetChildBaseFromEiaSteoClassifyById(classifyId int) (list []*BaseFromEiaSteoClassify, err error) {
	o := orm.NewOrmUsingDB("data")
	sql := `SELECT * FROM base_from_eia_steo_classify WHERE parent_id = ?`
	_, err = o.Raw(sql, classifyId).QueryRows(&list)
	return
}

func GetBaseFromEiaSteoClassifyByName(classifyName string) (item *BaseFromEiaSteoClassify, err error) {
	o := orm.NewOrmUsingDB("data")
	sql := `SELECT * FROM base_from_eia_steo_classify WHERE classify_name = ?`
	err = o.Raw(sql, classifyName).QueryRow(&item)
	return
}

// GetBaseFromEiaSteoIndexAll 获取所有的指标
func GetBaseFromEiaSteoIndexAll() (list []*BaseFromEiaSteoIndex, err error) {
	o := orm.NewOrmUsingDB("data")
	sql := `SELECT * FROM base_from_eia_steo_index`
	_, err = o.Raw(sql).QueryRows(&list)
	return
}

func lastDayOfMonth(year int, month time.Month) string {
	firstDay := time.Date(year, month, 1, 0, 0, 0, 0, time.Local)
	nextMonth := firstDay.AddDate(0, 1, 0)
	return nextMonth.Add(-time.Second * 1).Format(utils.FormatDate)
}

// HandleEiaSteoData 数据处理
func HandleEiaSteoData(dataList map[string]interface{}, indexInfo *BaseFromEiaSteoIndex) (err error) {
	o := orm.NewOrmUsingDB("data")
	to, err := o.Begin()
	if err != nil {
		return
	}
	defer func() {
		if err != nil {
			to.Rollback()
		} else {
			to.Commit()
		}
	}()
	var existList []*BaseFromEiaSteoData
	sql := `SELECT * FROM base_from_eia_steo_data WHERE index_code = ? order by data_time asc`
	_, err = to.Raw(sql, indexInfo.IndexCode).QueryRows(&existList)
	if err != nil {
		return err
	}

	existMap := make(map[string]*BaseFromEiaSteoData)
	for _, v := range existList {
		// 产品要求每个月的最后一天，为了兼容之前每个月第一天的数据
		dateTime := lastDayOfMonth(v.DataTime.Year(), v.DataTime.Month())
		existMap[dateTime] = v
	}

	// 待添加的数据列表
	addBaseFromEiaSteoDataList := make([]*BaseFromEiaSteoData, 0)
	endDate := indexInfo.EndDate
	startDate := indexInfo.StartDate

	for date, v := range dataList {
		// 实际数据
		val := v.(float64)
		valDecimal := decimal.NewFromFloat(val).Round(4)
		val, _ = valDecimal.Float64()

		currDate, tmpErr := time.ParseInLocation(utils.FormatMonthDateUnSpace, date, time.Local)
		if tmpErr != nil {
			err = tmpErr
			return
		}
		initDate := time.Date(2010, time.January, 1, 0, 0, 0, 0, time.Local)
		if currDate.Before(initDate) {
			continue
		}
		if currDate.After(endDate) {
			endDate = currDate
		}
		if currDate.Before(startDate) {
			startDate = currDate
		}
		// 产品要求每个月的最后一天
		date = lastDayOfMonth(currDate.Year(), currDate.Month())
		currDate, _ = time.ParseInLocation(utils.FormatDate, date, time.Local)
		// date = currDate.Format(utils.FormatDate)
		//fmt.Println("date:", date)

		tmpBaseFromEiaSteoData, ok := existMap[date]
		if !ok {
			addBaseFromEiaSteoDataList = append(addBaseFromEiaSteoDataList, &BaseFromEiaSteoData{
				//BaseFromEiaSteoDataId:  0,
				BaseFromEiaSteoIndexId: indexInfo.BaseFromEiaSteoIndexId,
				IndexCode:              indexInfo.IndexCode,
				DataTime:               currDate,
				Value:                  val,
				ModifyTime:             time.Now(),
				CreateTime:             time.Now(),
			})
		} else {
			nowVal := decimal.NewFromFloat(tmpBaseFromEiaSteoData.Value)
			//fmt.Println(date, "=========", nowVal.Equal(valDecimal))
			var updateCols []string
			if !nowVal.Equal(valDecimal) {
				tmpBaseFromEiaSteoData.Value = val
				tmpBaseFromEiaSteoData.ModifyTime = time.Now()
				updateCols = append(updateCols, "value")
				updateCols = append(updateCols, "modify_time")
			}
			if len(updateCols) > 0 {
				_, err = to.Update(tmpBaseFromEiaSteoData, updateCols...)
				if err != nil {
					return
				}
			}
		}
		delete(existMap, date)
	}

	lenAddBaseFromEiaSteoDataList := len(addBaseFromEiaSteoDataList)
	if lenAddBaseFromEiaSteoDataList > 0 {
		_, err = to.InsertMulti(lenAddBaseFromEiaSteoDataList, addBaseFromEiaSteoDataList)
	}

	existMapNum := len(existMap)
	if existMapNum > 0 {
		idList := make([]int, 0)
		for _, v := range existMap {
			idList = append(idList, v.BaseFromEiaSteoDataId)
		}

		sql = `DELETE FROM base_from_eia_steo_data WHERE base_from_eia_steo_data_id in(` + utils.GetOrmInReplace(existMapNum) + `) `
		_, err = to.Raw(sql, idList).Exec()
		if err != nil {
			return err
		}
	}

	// 更新指标信息
	{
		updateParams := make([]string, 0)
		indexInfo.ModifyTime = time.Now()
		updateParams = append(updateParams, "ModifyTime")
		if endDate.After(indexInfo.EndDate) {
			indexInfo.EndDate = endDate
			updateParams = append(updateParams, "EndDate")
		}
		if startDate.Before(indexInfo.StartDate) {
			indexInfo.StartDate = startDate
			updateParams = append(updateParams, "StartDate")
		}
		_, err = to.Update(indexInfo, updateParams...)
		if err != nil {
			return
		}
	}
	return
}