AnthoLume/database/manager.go

package database

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"embed"
	_ "embed"
	"errors"
	"fmt"
	"path/filepath"
	"time"

	"github.com/pressly/goose/v3"
	log "github.com/sirupsen/logrus"
	sqlite "modernc.org/sqlite"
	"reichard.io/antholume/config"
	_ "reichard.io/antholume/database/migrations"
)

type DBManager struct {
	DB      *sql.DB
	Ctx     context.Context
	Queries *Queries
	cfg     *config.Config
}

//go:embed schema.sql
var ddl string

//go:embed user_streaks.sql
var user_streaks string

//go:embed document_user_statistics.sql
var document_user_statistics string

//go:embed migrations/*
var migrations embed.FS

// Register scalar sqlite function on init
func init() {
	sqlite.MustRegisterFunction("LOCAL_TIME", &sqlite.FunctionImpl{
		NArgs:         2,
		Deterministic: true,
		Scalar:        localTime,
	})
	sqlite.MustRegisterFunction("LOCAL_DATE", &sqlite.FunctionImpl{
		NArgs:         2,
		Deterministic: true,
		Scalar:        localDate,
	})
}

// NewMgr Returns an initialized manager
func NewMgr(c *config.Config) *DBManager {
	// Create Manager
	dbm := &DBManager{
		Ctx: context.Background(),
		cfg: c,
	}

	if err := dbm.init(); err != nil {
		log.Panic("Unable to init DB")
	}

	return dbm
}

// init loads the DB manager
func (dbm *DBManager) init() error {
	// Build DB Location
	var dbLocation string
	switch dbm.cfg.DBType {
	case "sqlite":
		dbLocation = filepath.Join(dbm.cfg.ConfigPath, fmt.Sprintf("%s.db", dbm.cfg.DBName))
	case "memory":
		dbLocation = ":memory:"
	default:
		return fmt.Errorf("unsupported database")
	}

	var err error
	dbm.DB, err = sql.Open("sqlite", dbLocation)
	if err != nil {
		log.Panicf("Unable to open DB: %v", err)
		return err
	}

	// Single open connection
	dbm.DB.SetMaxOpenConns(1)

	// Check if DB is new
	isNew, err := isEmpty(dbm.DB)
	if err != nil {
		log.Panicf("Unable to determine db info: %v", err)
		return err
	}

	// Init SQLc
	dbm.Queries = New(dbm.DB)

	// Execute schema
	if _, err := dbm.DB.Exec(ddl, nil); err != nil {
		log.Panicf("Error executing schema: %v", err)
		return err
	}

	// Perform migrations
	err = dbm.performMigrations(isNew)
	if err != nil && err != goose.ErrNoMigrationFiles {
		log.Panicf("Error running DB migrations: %v", err)
		return err
	}

	// Update settings
	err = dbm.updateSettings()
	if err != nil {
		log.Panicf("Error running DB settings update: %v", err)
		return err
	}

	// Cache tables
	if err := dbm.CacheTempTables(); err != nil {
		log.Warn("Refreshing temp table cache failed: ", err)
	}

	return nil
}

// Reload closes the DB & reinits
func (dbm *DBManager) Reload() error {
	// Close handle
	err := dbm.DB.Close()
	if err != nil {
		return err
	}

	// Reinit DB
	if err := dbm.init(); err != nil {
		return err
	}

	return nil
}

// CacheTempTables clears existing statistics and recalculates
func (dbm *DBManager) CacheTempTables() error {
	start := time.Now()
	if _, err := dbm.DB.ExecContext(dbm.Ctx, user_streaks); err != nil {
		return err
	}
	log.Debug("Cached 'user_streaks' in: ", time.Since(start))

	start = time.Now()
	if _, err := dbm.DB.ExecContext(dbm.Ctx, document_user_statistics); err != nil {
		return err
	}
	log.Debug("Cached 'document_user_statistics' in: ", time.Since(start))

	return nil
}

// updateSettings ensures that we're enforcing foreign keys and enable journal
// mode.
func (dbm *DBManager) updateSettings() error {
	// Set SQLite PRAGMA Settings
	pragmaQuery := `
		  PRAGMA foreign_keys = ON;
		  PRAGMA journal_mode = WAL;
		`
	if _, err := dbm.DB.Exec(pragmaQuery, nil); err != nil {
		log.Errorf("Error executing pragma: %v", err)
		return err
	}

	// Update Antholume Version in DB
	if _, err := dbm.Queries.UpdateSettings(dbm.Ctx, UpdateSettingsParams{
		Name:  "version",
		Value: dbm.cfg.Version,
	}); err != nil {
		log.Errorf("Error updating DB settings: %v", err)
		return err
	}

	return nil
}

// performMigrations runs all migrations
func (dbm *DBManager) performMigrations(isNew bool) error {
	// Create context
	ctx := context.WithValue(context.Background(), "isNew", isNew) // nolint

	// Set DB migration
	goose.SetBaseFS(migrations)

	// Run migrations
	goose.SetLogger(log.StandardLogger())
	if err := goose.SetDialect("sqlite"); err != nil {
		return err
	}

	return goose.UpContext(ctx, dbm.DB, "migrations")
}

// isEmpty determines whether the database is empty
func isEmpty(db *sql.DB) (bool, error) {
	var tableCount int
	err := db.QueryRow("SELECT COUNT(*) FROM sqlite_master WHERE type='table';").Scan(&tableCount)
	if err != nil {
		return false, err
	}
	return tableCount == 0, nil
}

// localTime is a custom SQL function that is registered as LOCAL_TIME in the init function
func localTime(ctx *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
	timeStr, ok := args[0].(string)
	if !ok {
		return nil, errors.New("both arguments to TZTime must be strings")
	}

	timeZoneStr, ok := args[1].(string)
	if !ok {
		return nil, errors.New("both arguments to TZTime must be strings")
	}

	timeZone, err := time.LoadLocation(timeZoneStr)
	if err != nil {
		return nil, errors.New("unable to parse timezone")
	}

	formattedTime, err := time.ParseInLocation(time.RFC3339, timeStr, time.UTC)
	if err != nil {
		return nil, errors.New("unable to parse time")
	}

	return formattedTime.In(timeZone).Format(time.RFC3339), nil
}

// localDate is a custom SQL function that is registered as LOCAL_DATE in the init function
func localDate(ctx *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
	timeStr, ok := args[0].(string)
	if !ok {
		return nil, errors.New("both arguments to TZTime must be strings")
	}

	timeZoneStr, ok := args[1].(string)
	if !ok {
		return nil, errors.New("both arguments to TZTime must be strings")
	}

	timeZone, err := time.LoadLocation(timeZoneStr)
	if err != nil {
		return nil, errors.New("unable to parse timezone")
	}

	formattedTime, err := time.ParseInLocation(time.RFC3339, timeStr, time.UTC)
	if err != nil {
		return nil, errors.New("unable to parse time")
	}

	return formattedTime.In(timeZone).Format("2006-01-02"), nil
}