twitch-bot-streak/pkg/database/db.go

package database

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/glebarez/sqlite"
	"github.com/pkg/errors"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type (
	DB struct {
		db *gorm.DB
	}
)

// New creates a new DB with the given driver and database
func New(driverName, connString string) (d *DB, err error) {
	var innerDB gorm.Dialector

	switch driverName {
	case "mysql":
		innerDB = mysql.Open(connString)

	case "postgres":
		innerDB = postgres.Open(connString)

	case "sqlite":
		innerDB = sqlite.Open(connString)

	default:
		return nil, errors.Errorf("unknown database driver %s", driverName)
	}

	db, err := gorm.Open(innerDB, &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
		Logger: logger.New(log.New(os.Stderr, "\r\n", log.LstdFlags), logger.Config{
			SlowThreshold:             time.Second,
			LogLevel:                  logger.Warn,
			IgnoreRecordNotFoundError: true,
			Colorful:                  false,
		}),
	})
	if err != nil {
		return nil, errors.Wrap(err, "connecting database")
	}

	d = &DB{
		db: db,
	}

	if err = db.AutoMigrate(
		&StreakMeta{},
		&StreakUser{},
	); err != nil {
		return nil, fmt.Errorf("applying schema: %w", err)
	}

	return d, nil
}

// Close closes the underlying database
func (d DB) Close() error {
	sdb, err := d.db.DB()
	if err != nil {
		return fmt.Errorf("getting underlying database: %w", err)
	}

	if err = sdb.Close(); err != nil {
		return fmt.Errorf("closing db: %w", err)
	}

	return nil
}