spider_creater/backend/db/db.go

package db

import (
	"encoding/json"
	"errors"
	"github.com/boltdb/bolt"
	qu "jygit.jydev.jianyu360.cn/data_processing/common_utils"
	be "spider_creator/backend"
)

// 在新增表的时候，需要在这里增加实体清单
type Entity interface {
	be.SpiderConfig | be.Job
}

var (
	// 在新增表的时候，需要在这里增加表名
	DB_TABLES = []string{"myBucket", "jobs", "spiderconfig", "systemconfig"}
	Db        *SpiderDb
)

type (
	//SpiderDB 爬虫库，这里模拟真实数据库
	SpiderDb struct {
		db  *bolt.DB
		enf be.EventNotifyFace
	}
)

// NewSpiderDb
func NewSpiderDb(dbfile string, enf be.EventNotifyFace) *SpiderDb {
	db, err := bolt.Open(dbfile, 0600, nil)
	if err != nil {
		qu.Debug("db error", err.Error())
	}
	for _, t := range DB_TABLES {
		err = db.Update(func(tx *bolt.Tx) error {
			_, err := tx.CreateBucketIfNotExists([]byte(t))
			return err
		})
		if err != nil {
			qu.Debug("db error", err.Error())
		}
	}
	return &SpiderDb{
		db, enf,
	}
}

// Close
func (s *SpiderDb) Close() {
	s.db.Close()
}

// 支持泛型的通用方法
// LoadEntity
func LoadEntity[T any](table, key string) (*T, error) {
	var ret *T = new(T)
	err := Db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(table))
		value := bucket.Get([]byte(key))
		if value != nil && len(value) > 0 {
			_ = json.Unmarshal(value, ret)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return ret, nil
}

// LoadEntities
func LoadEntities[T1 Entity](table string) ([]*T1, error) {
	ret := make([]*T1, 0)
	// 开始读取事务
	err := Db.db.View(func(tx *bolt.Tx) error {
		// 遍历数据库中的所有桶
		bucket := tx.Bucket([]byte(table))
		if bucket == nil {
			return errors.New("桶不存在")
		}
		// 遍历桶中的所有键/值对
		return bucket.ForEach(func(k, v []byte) error {
			var sf = new(T1)
			json.Unmarshal(v, sf)
			if sf != nil {
				ret = append(ret, sf)
			}
			return nil
		})
	})
	if err != nil {
		return ret, err
	}

	return ret, nil
}

// DeleteEntity
func DeleteEntity[T any](table, key string) error {
	return Db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(table))
		return bucket.Delete([]byte(key))

	})
}

// SaveEntity
func SaveEntity[T any](table, key string, obj *T) error {
	//加载原始数据
	value, _ := json.Marshal(obj)
	return Db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(table))
		err := bucket.Put([]byte(key), value)
		return err
	})
}