data_workshop/base.go

package main

import (
	"bytes"
	"compress/flate"
	"encoding/gob"
	"log"
	"math/rand"
	"os"
	"reflect"
)

type KeyKind struct {
	Key       string
	FieldName string
}

func NewKeyKind(key string, fieldname string) KeyKind {
	return KeyKind{Key: key, FieldName: fieldname}
}

type KeyList struct {
	Keys []interface{}
}

func (kl *KeyList) GetKeys() []interface{} {
	return kl.Keys
}

func (kl *KeyList) AppendKey(key interface{}) {
	kl.Keys = append(kl.Keys, key)
}

func (kl *KeyList) GetLength() int {
	return len(kl.Keys)
}

// CheckErrorPanic 查错误 存在就panic操作
func CheckErrorPanic(err error) {
	if err != nil {
		log.Panic(err)
	}
}

// OpenDB 打开rocksdb
// func OpenDB() *gorocksdb.DB {
// 	bbto := gorocksdb.NewDefaultBlockBasedTableOptions()

// 	bbto.SetBlockCache(gorocksdb.NewLRUCache(2 << 30))
// 	bbto.SetCacheIndexAndFilterBlocksWithHighPriority(true)

// 	opts := gorocksdb.NewDefaultOptions()

// 	bbto.SetFilterPolicy(gorocksdb.NewBloomFilter(16))

// 	opts.SetBlockBasedTableFactory(bbto)
// 	opts.SetCreateIfMissing(true)
// 	opts.SetCreateIfMissingColumnFamilies(true)
// 	opts.SetCompression(gorocksdb.LZ4Compression)

// 	db, err := gorocksdb.OpenDb(opts, ".rocksdb")
// 	CheckErrorPanic(err)
// 	return db
// }

// var db = OpenDB()
// var wopts = gorocksdb.NewDefaultWriteOptions()
// var ropts = gorocksdb.NewDefaultReadOptions()

// IKeyList key list interface
type IKeyList interface {
	AppendKey(key interface{})
	GetKeys() []interface{}
	GetLength() int
}

// SaveData gob序列化存数据
func SaveData(fname string, v interface{}) {
	var keybuf = &bytes.Buffer{}

	enc := gob.NewEncoder(keybuf)
	enc.Encode(v)

	f, err := os.OpenFile(fname, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.ModePerm)
	CheckErrorPanic(err)

	gw, err := flate.NewWriter(f, 7)
	CheckErrorPanic(err)
	defer gw.Close()

	gw.Write(keybuf.Bytes())
}

// SaveGob 用于解析rocksdb源数据, 爬取数据存储在rocksdb后取出 序列化存储
// func SaveGob(fname string, datatype reflect.Type, kfs ...KeyKind) {

// 	// cl := &CountryList{}
// 	kl := &KeyList{}

// 	i := 0
// 	iter := db.NewIterator(ropts)
// 	bkey := []byte(kfs[0].Key)
// 	iter.Seek(bkey)

// 	for ; iter.ValidForPrefix(bkey); iter.Next() {
// 		i++
// 		first := iter.Value().Data()
// 		data := reflect.New(datatype).Elem()

// 		data.FieldByName(kfs[0].FieldName).Set(reflect.ValueOf(first))

// 		// country := &Country{}
// 		// country.Name = name
// 		for n := 1; n < len(kfs); n++ {
// 			v, err := db.Get(ropts, []byte(fmt.Sprintf(kfs[n].Key, first)))
// 			if err != nil {
// 				panic(err)
// 			}
// 			fv := data.FieldByName(kfs[n].FieldName)
// 			fv.Set(reflect.ValueOf(v.Data()))
// 		}
// 		kl.AppendKey(data.Interface())
// 	}
// 	SaveData(fname, kl)
// }

// LoadGob load gob from file
func LoadGob(fname string, v IKeyList) {

	f, err := os.Open(fname)
	log.Println("load data: ", fname)
	CheckErrorPanic(err)
	reader := flate.NewReader(f)
	dec := gob.NewDecoder(reader)
	CheckErrorPanic(dec.Decode(v))

}

// GetRandomKey get ikeylist key by random
func GetRandomKey(v IKeyList) interface{} {
	ret := v.GetKeys()[rand.Intn(v.GetLength())]
	return ret
}

// GetRandomKeyBySlice get ikeylist key by random
func GetRandomKeyBySlice(v interface{}) interface{} {
	sli := reflect.ValueOf(v)
	return sli.Index(rand.Intn(sli.Len())).Interface()
}