rushlink/internal/db/db.go

package db

import (
	"fmt"
	"log"
	"time"

	"github.com/pkg/errors"

	bolt "go.etcd.io/bbolt"

	gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/pkg/gobmarsh"
)

// Database is the main rushlink database type.
//
// Open a database using DB.Open() and close it in the end using DB.Close().
// Only one instance of DB should exist in a program at any moment.
type Database struct {
	Bolt *bolt.DB
}

// CurrentMigrateVersion holds the current "migrate version".
//
// If we alter the database format, we bump this number and write a new
// database migration in migrate().
const CurrentMigrateVersion = 2

// BucketConf holds the name for the "configuration" bucket.
//
// This bucket holds the database version, secret site-wide keys, etc.
const BucketConf = "conf"

// BucketPastes holds the name for the pastes bucket.
const BucketPastes = "pastes"

// BucketFileUpload holds the name for the file-upload bucket.
const BucketFileUpload = "fileUpload"

// KeyMigrateVersion stores the current migration version. If this value is less than
// CurrentMigrateVersion, the database has to be migrated.
const KeyMigrateVersion = "migrate_version"

// OpenDB opens a database file located at path.
func OpenDB(path string) (*Database, error) {
	if path == "" {
		return nil, errors.New("database not set")
	}

	db, err := bolt.Open(path, 0666, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, errors.Wrapf(err, "failed to open database at '%v'", path)
	}
	if err := db.Update(migrate); err != nil {
		return nil, err
	}
	return &Database{db}, nil
}

// Close the bolt database
func (db *Database) Close() error {
	if db == nil {
		panic("no open database")
	}
	return db.Close()
}

// Initialize and migrate the database to the current version
func migrate(tx *bolt.Tx) error {
	dbVersion, err := dbVersion(tx)
	if err != nil {
		return err
	}

	// Migrate the database to version 1
	if dbVersion < 1 {
		log.Println("migrating database to version 1")
		// Create conf bucket
		_, err := tx.CreateBucket([]byte(BucketConf))
		if err != nil {
			return err
		}
		// Create paste bucket
		_, err = tx.CreateBucket([]byte(BucketPastes))
		if err != nil {
			return err
		}
		// Update the version number
		if err := setDBVersion(tx, 1); err != nil {
			return err
		}
	}

	if dbVersion < 2 {
		log.Println("migrating database to version 2")
		// Create fileUpload bucket
		_, err := tx.CreateBucket([]byte(BucketFileUpload))
		if err != nil {
			return err
		}
		// Update the version number
		if err := setDBVersion(tx, 2); err != nil {
			return err
		}
	}

	return nil
}

// Get the current migrate version from the database
func dbVersion(tx *bolt.Tx) (int, error) {
	conf := tx.Bucket([]byte(BucketConf))
	if conf == nil {
		return 0, nil
	}
	dbVersionBytes := conf.Get([]byte(KeyMigrateVersion))
	if dbVersionBytes == nil {
		return 0, nil
	}

	// Version was already stored
	var dbVersion int
	if err := gobmarsh.Unmarshal(dbVersionBytes, &dbVersion); err != nil {
		return 0, err
	}
	if dbVersion == 0 {
		return 0, fmt.Errorf("database version is invalid (%v)", dbVersion)
	}
	if dbVersion > CurrentMigrateVersion {
		return 0, fmt.Errorf("database version is too recent (%v > %v)", dbVersion, CurrentMigrateVersion)
	}
	return dbVersion, nil
}

// Update the current migrate version in the database
func setDBVersion(tx *bolt.Tx, version int) error {
	conf, err := tx.CreateBucketIfNotExists([]byte(BucketConf))
	if err != nil {
		return err
	}

	versionBytes, err := gobmarsh.Marshal(version)
	if err != nil {
		return err
	}
	return conf.Put([]byte(KeyMigrateVersion), versionBytes)
}