tunasync/manager/db.go

package manager

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/boltdb/bolt"
	"github.com/dgraph-io/badger/v2"
	"github.com/go-redis/redis/v8"

	. "github.com/tuna/tunasync/internal"
)

type dbAdapter interface {
	Init() error
	ListWorkers() ([]WorkerStatus, error)
	GetWorker(workerID string) (WorkerStatus, error)
	DeleteWorker(workerID string) error
	CreateWorker(w WorkerStatus) (WorkerStatus, error)
	RefreshWorker(workerID string) (WorkerStatus, error)
	UpdateMirrorStatus(workerID, mirrorID string, status MirrorStatus) (MirrorStatus, error)
	GetMirrorStatus(workerID, mirrorID string) (MirrorStatus, error)
	ListMirrorStatus(workerID string) ([]MirrorStatus, error)
	ListAllMirrorStatus() ([]MirrorStatus, error)
	FlushDisabledJobs() error
	Close() error
}

// interface for a kv database
type kvAdapter interface {
	InitBucket(bucket string) error
	Get(bucket string, key string) ([]byte, error)
	GetAll(bucket string) (map[string][]byte, error)
	Put(bucket string, key string, value []byte) error
	Delete(bucket string, key string) error
	Close() error
}

const (
	_workerBucketKey = "workers"
	_statusBucketKey = "mirror_status"
)

func makeDBAdapter(dbType string, dbFile string) (dbAdapter, error) {
	if dbType == "bolt" {
		innerDB, err := bolt.Open(dbFile, 0600, &bolt.Options{
			Timeout: 5 * time.Second,
		})
		if err != nil {
			return nil, err
		}
		db := boltAdapter{
			db: innerDB,
		}
		kv := kvDBAdapter{
			db: &db,
		}
		err = kv.Init()
		return &kv, err
	} else if dbType == "redis" {
		opt, err := redis.ParseURL(dbFile)
		if err != nil {
			return nil, fmt.Errorf("bad redis url: %s", err)
		}
		innerDB := redis.NewClient(opt)
		db := redisAdapter{
			db: innerDB,
		}
		kv := kvDBAdapter{
			db: &db,
		}
		err = kv.Init()
		return &kv, err
	} else if dbType == "badger" {
		innerDB, err := badger.Open(badger.DefaultOptions(dbFile))
		if err != nil {
			return nil, err
		}
		db := badgerAdapter{
			db: innerDB,
		}
		kv := kvDBAdapter{
			db: &db,
		}
		err = kv.Init()
		return &kv, err
	}
	// unsupported db-type
	return nil, fmt.Errorf("unsupported db-type: %s", dbType)
}

// use the underlying kv database to store data
type kvDBAdapter struct {
	db kvAdapter
}

func (b *kvDBAdapter) Init() error {
	err := b.db.InitBucket(_workerBucketKey)
	if err != nil {
		return fmt.Errorf("create bucket %s error: %s", _workerBucketKey, err.Error())
	}
	err = b.db.InitBucket(_statusBucketKey)
	if err != nil {
		return fmt.Errorf("create bucket %s error: %s", _workerBucketKey, err.Error())
	}
	return err
}

func (b *kvDBAdapter) ListWorkers() (ws []WorkerStatus, err error) {
	var workers map[string][]byte
	workers, err = b.db.GetAll(_workerBucketKey)

	var w WorkerStatus
	for _, v := range workers {
		jsonErr := json.Unmarshal(v, &w)
		if jsonErr != nil {
			err = fmt.Errorf("%s; %s", err.Error(), jsonErr)
			continue
		}
		ws = append(ws, w)
	}
	return
}

func (b *kvDBAdapter) GetWorker(workerID string) (w WorkerStatus, err error) {
	var v []byte
	v, err = b.db.Get(_workerBucketKey, workerID)
	if v == nil {
		err = fmt.Errorf("invalid workerID %s", workerID)
	} else {
		err = json.Unmarshal(v, &w)
	}
	return
}

func (b *kvDBAdapter) DeleteWorker(workerID string) error {
	v, _ := b.db.Get(_workerBucketKey, workerID)
	if v == nil {
		return fmt.Errorf("invalid workerID %s", workerID)
	}
	return b.db.Delete(_workerBucketKey, workerID)
}

func (b *kvDBAdapter) CreateWorker(w WorkerStatus) (WorkerStatus, error) {
	v, err := json.Marshal(w)
	if err == nil {
		err = b.db.Put(_workerBucketKey, w.ID, v)
	}
	return w, err
}

func (b *kvDBAdapter) RefreshWorker(workerID string) (w WorkerStatus, err error) {
	w, err = b.GetWorker(workerID)
	if err == nil {
		w.LastOnline = time.Now()
		w, err = b.CreateWorker(w)
	}
	return w, err
}

func (b *kvDBAdapter) UpdateMirrorStatus(workerID, mirrorID string, status MirrorStatus) (MirrorStatus, error) {
	id := mirrorID + "/" + workerID
	v, err := json.Marshal(status)
	if err == nil {
		err = b.db.Put(_statusBucketKey, id, v)
	}
	return status, err
}

func (b *kvDBAdapter) GetMirrorStatus(workerID, mirrorID string) (m MirrorStatus, err error) {
	id := mirrorID + "/" + workerID
	var v []byte
	v, err = b.db.Get(_statusBucketKey, id)
	if v == nil {
		err = fmt.Errorf("no mirror '%s' exists in worker '%s'", mirrorID, workerID)
	} else if err == nil {
		err = json.Unmarshal(v, &m)
	}
	return
}

func (b *kvDBAdapter) ListMirrorStatus(workerID string) (ms []MirrorStatus, err error) {
	var vals map[string][]byte
	vals, err = b.db.GetAll(_statusBucketKey)
	if err != nil {
		return
	}

	for k, v := range vals {
		if wID := strings.Split(k, "/")[1]; wID == workerID {
			var m MirrorStatus
			jsonErr := json.Unmarshal(v, &m)
			if jsonErr != nil {
				err = fmt.Errorf("%s; %s", err.Error(), jsonErr)
				continue
			}
			ms = append(ms, m)
		}
	}
	return
}

func (b *kvDBAdapter) ListAllMirrorStatus() (ms []MirrorStatus, err error) {
	var vals map[string][]byte
	vals, err = b.db.GetAll(_statusBucketKey)
	if err != nil {
		return
	}

	for _, v := range vals {
		var m MirrorStatus
		jsonErr := json.Unmarshal(v, &m)
		if jsonErr != nil {
			err = fmt.Errorf("%s; %s", err.Error(), jsonErr)
			continue
		}
		ms = append(ms, m)
	}
	return
}

func (b *kvDBAdapter) FlushDisabledJobs() (err error) {
	var vals map[string][]byte
	vals, err = b.db.GetAll(_statusBucketKey)
	if err != nil {
		return
	}

	for k, v := range vals {
		var m MirrorStatus
		jsonErr := json.Unmarshal(v, &m)
		if jsonErr != nil {
			err = fmt.Errorf("%s; %s", err.Error(), jsonErr)
			continue
		}
		if m.Status == Disabled || len(m.Name) == 0 {
			deleteErr := b.db.Delete(_statusBucketKey, k)
			if deleteErr != nil {
				err = fmt.Errorf("%s; %s", err.Error(), deleteErr)
			}
		}
	}
	return
}

func (b *kvDBAdapter) Close() error {
	if b.db != nil {
		return b.db.Close()
	}
	return nil
}