memos/store/sqlite.go

package store

import (
	"database/sql"
	"embed"
	"errors"
	"fmt"
	"io/fs"
	"memos/common"
	"os"
	"sort"

	_ "github.com/mattn/go-sqlite3"
)

//go:embed migration
var migrationFS embed.FS

//go:embed seed
var seedFS embed.FS

type DB struct {
	Db *sql.DB
	// datasource name
	DSN string
	// mode should be prod or dev
	mode string
}

// NewDB returns a new instance of DB associated with the given datasource name.
func NewDB(profile *common.Profile) *DB {
	db := &DB{
		DSN:  profile.DSN,
		mode: profile.Mode,
	}
	return db
}

func (db *DB) Open() (err error) {
	// Ensure a DSN is set before attempting to open the database.
	if db.DSN == "" {
		return fmt.Errorf("dsn required")
	}

	// Connect to the database.
	sqlDB, err := sql.Open("sqlite3", db.DSN)
	if err != nil {
		return fmt.Errorf("failed to open db with dsn: %s, err: %w", db.DSN, err)
	}

	db.Db = sqlDB
	// If db file not exists, we should migrate and seed the database.
	if _, err := os.Stat(db.DSN); errors.Is(err, os.ErrNotExist) {
		if err := db.migrate(); err != nil {
			return fmt.Errorf("failed to migrate: %w", err)
		}
		// If mode is dev, then seed the database.
		if db.mode == "dev" {
			if err := db.seed(); err != nil {
				return fmt.Errorf("failed to seed: %w", err)
			}
		}
	} else {
		// If db file exists and mode is dev, we should migrate and seed the database.
		if db.mode == "dev" {
			if err := db.migrate(); err != nil {
				return fmt.Errorf("failed to migrate: %w", err)
			}
			if err := db.seed(); err != nil {
				return fmt.Errorf("failed to seed: %w", err)
			}
		}
	}

	return err
}

func (db *DB) migrate() error {
	filenames, err := fs.Glob(migrationFS, fmt.Sprintf("%s/*.sql", "migration"))
	if err != nil {
		return err
	}

	sort.Strings(filenames)

	// Loop over all migration files and execute them in order.
	for _, filename := range filenames {
		if err := db.executeFile(migrationFS, filename); err != nil {
			return fmt.Errorf("migrate error: name=%q err=%w", filename, err)
		}
	}
	return nil
}

func (db *DB) seed() error {
	filenames, err := fs.Glob(seedFS, fmt.Sprintf("%s/*.sql", "seed"))
	if err != nil {
		return err
	}

	sort.Strings(filenames)

	// Loop over all seed files and execute them in order.
	for _, filename := range filenames {
		if err := db.executeFile(seedFS, filename); err != nil {
			return fmt.Errorf("seed error: name=%q err=%w", filename, err)
		}
	}
	return nil
}

// executeFile runs a single seed file within a transaction.
func (db *DB) executeFile(FS embed.FS, name string) error {
	tx, err := db.Db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// Read and execute SQL file.
	if buf, err := fs.ReadFile(FS, name); err != nil {
		return err
	} else if _, err := tx.Exec(string(buf)); err != nil {
		return err
	}

	return tx.Commit()
}

func FormatError(err error) error {
	if err == nil {
		return nil
	}

	switch err {
	case sql.ErrNoRows:
		return errors.New("data not found")
	default:
		return err
	}
}