data_workshop/base.go
2020-06-03 14:27:09 +08:00

150 lines
3.3 KiB
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()
}