go_dreamfactory/cmd/v2/service/dbServer.go

package service

import (
	"encoding/json"
	"fmt"
	"go_dreamfactory/cmd/v2/lib/common"
	"go_dreamfactory/cmd/v2/model"
	"log"
	"time"

	"github.com/boltdb/bolt"
	"github.com/sirupsen/logrus"
)

type DbService interface {
	// 保存鲁班工具配置
	SaveLubanConf(conf *model.GenTool) error
	// 获取鲁班工具配置
	GetLubanConf(key string) *model.GenTool
	// 保存SSH连接信息
	SaveSSHConf(conf *model.SSHModel) error
	//获取终端配置
	GetSSHConf(key string) *model.SSHModel
	// 保存pb配置
	SavePbConf(conf *model.PbConfModel) error
	// 获取pb配置
	GetPbConf() *model.PbConfModel
}

type DbServiceImpl struct {
}

var dbservice DbService

func NewDbService() DbService {
	dbservice = &DbServiceImpl{}
	return dbservice
}

func GetDbService() DbService {
	return dbservice
}

func (this *DbServiceImpl) save(key string, model interface{}) error {
	boltDb = GetBoltDb()
	defer boltDb.Close()
	if err := boltDb.Update(func(tx *bolt.Tx) error {
		b, err2 := json.Marshal(model)
		if err2 != nil {
			return err
		}
		c := tx.Bucket([]byte(common.BOLTDB_BUCKETNAME))
		if c == nil {
			return fmt.Errorf("Bucket %s not found!", common.BOLTDB_BUCKETNAME)
		}
		return c.Put([]byte(key), b)
	}); err != nil {
		logrus.Errorf("save err: %v", err)
		return err
	}
	return nil
}

func (this *DbServiceImpl) SaveLubanConf(conf *model.GenTool) error {
	return this.save(common.BUCKET_LUBANCONF, conf)
}

func (this *DbServiceImpl) SavePbConf(conf *model.PbConfModel) error {
	return this.save(common.BUCKET_PBCONF, conf)
}

func (this *DbServiceImpl) GetLubanConf(key string) *model.GenTool {
	boltDb = GetBoltDb()
	defer boltDb.Close()
	model := &model.GenTool{}
	if err = boltDb.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(common.BOLTDB_BUCKETNAME))
		if b == nil {
			return fmt.Errorf("Bucket %s not found!", common.BOLTDB_BUCKETNAME)
		}
		val := b.Get([]byte(key))
		if err = json.Unmarshal(val, model); err != nil {
			logrus.Errorf("get gen conf err:%v", err)
			return err
		}
		return nil
	}); err != nil {
		logrus.Errorf("get db conf err: %v", err)
		return nil
	}

	return model
}

func (this *DbServiceImpl) SaveSSHConf(conf *model.SSHModel) error {
	return this.save(common.BUCKET_SSHCONF, conf)
}

func (this *DbServiceImpl) GetSSHConf(key string) *model.SSHModel {
	boltDb = GetBoltDb()
	defer boltDb.Close()
	model := &model.SSHModel{}
	if err = boltDb.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(common.BOLTDB_BUCKETNAME))
		if b == nil {
			return fmt.Errorf("Bucket %s not found!", common.BOLTDB_BUCKETNAME)
		}
		val := b.Get([]byte(key))
		if err = json.Unmarshal(val, model); err != nil {
			logrus.Errorf("get gen conf err:%v", err)
			return err
		}
		return nil
	}); err != nil {
		logrus.Errorf("get db conf err: %v", err)
		return nil
	}

	return model
}

func (this *DbServiceImpl) GetPbConf() *model.PbConfModel {
	boltDb = GetBoltDb()
	defer boltDb.Close()
	model := &model.PbConfModel{}
	if err = boltDb.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(common.BOLTDB_BUCKETNAME))
		if b == nil {
			return fmt.Errorf("Bucket %s not found!", common.BOLTDB_BUCKETNAME)
		}
		val := b.Get([]byte(common.BUCKET_PBCONF))
		if err = json.Unmarshal(val, model); err != nil {
			logrus.Errorf("get gen conf err:%v", err)
			return err
		}
		return nil
	}); err != nil {
		logrus.Errorf("get db conf err: %v", err)
		return nil
	}

	return model
}

var (
	boltDb *bolt.DB
	bucket *bolt.Bucket
	err    error
)

func GetBoltDb() *bolt.DB {
	boltDb, err = bolt.Open(common.BOLTDB_NAME, 0600, &bolt.Options{Timeout: 5 * time.Second})
	if err != nil {
		log.Fatal(err)
	}

	boltDb.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(common.BOLTDB_BUCKETNAME))
		if b != nil {
			bucket = b
		} else {
			bucket, err = tx.CreateBucket([]byte(common.BOLTDB_BUCKETNAME))
			if err != nil {
				return fmt.Errorf("create bucket: %s", err)
			}
		}
		return nil
	})
	return boltDb
}