chore: move migration and seed code into driver (#2294)

Move migration and seed code into driver

cmd/memos.go

@@ -16,7 +16,6 @@ import (
 	"github.com/usememos/memos/server"
 	_profile "github.com/usememos/memos/server/profile"
 	"github.com/usememos/memos/store"
-	"github.com/usememos/memos/store/db"
 	"github.com/usememos/memos/store/sqlite"
 )
 
@@ -43,19 +42,18 @@ var (
 		Short: `An open-source, self-hosted memo hub with knowledge management and social networking.`,
 		Run: func(_cmd *cobra.Command, _args []string) {
 			ctx, cancel := context.WithCancel(context.Background())
-			db := db.NewDB(profile)
-			if err := db.Open(); err != nil {
+			driver, err := sqlite.NewDriver(profile)
+			if err != nil {
 				cancel()
-				log.Error("failed to open db", zap.Error(err))
+				log.Error("failed to create db driver", zap.Error(err))
 				return
 			}
-			if err := db.Migrate(ctx); err != nil {
+			if err := driver.Migrate(ctx); err != nil {
 				cancel()
 				log.Error("failed to migrate db", zap.Error(err))
 				return
 			}
 
-			driver := sqlite.NewDriver(db.DBInstance)
 			store := store.New(driver, profile)
 			s, err := server.NewServer(ctx, profile, store)
 			if err != nil {

cmd/mvrss.go

@@ -9,7 +9,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/usememos/memos/store"
-	"github.com/usememos/memos/store/db"
 	"github.com/usememos/memos/store/sqlite"
 )
 
@@ -40,17 +39,16 @@ var (
 				return
 			}
 
-			db := db.NewDB(profile)
-			if err := db.Open(); err != nil {
-				fmt.Printf("failed to open db, error: %+v\n", err)
+			driver, err := sqlite.NewDriver(profile)
+			if err != nil {
+				fmt.Printf("failed to create db driver, error: %+v\n", err)
 				return
 			}
-			if err := db.Migrate(ctx); err != nil {
+			if err := driver.Migrate(ctx); err != nil {
 				fmt.Printf("failed to migrate db, error: %+v\n", err)
 				return
 			}
 
-			driver := sqlite.NewDriver(db.DBInstance)
 			s := store.New(driver, profile)
 			resources, err := s.ListResources(ctx, &store.FindResource{})
 			if err != nil {

cmd/setup.go

@@ -11,7 +11,6 @@ import (
 
 	"github.com/usememos/memos/common/util"
 	"github.com/usememos/memos/store"
-	"github.com/usememos/memos/store/db"
 	"github.com/usememos/memos/store/sqlite"
 )
 
@@ -37,17 +36,16 @@ var (
 				return
 			}
 
-			db := db.NewDB(profile)
-			if err := db.Open(); err != nil {
-				fmt.Printf("failed to open db, error: %+v\n", err)
+			driver, err := sqlite.NewDriver(profile)
+			if err != nil {
+				fmt.Printf("failed to create db driver, error: %+v\n", err)
 				return
 			}
-			if err := db.Migrate(ctx); err != nil {
+			if err := driver.Migrate(ctx); err != nil {
 				fmt.Printf("failed to migrate db, error: %+v\n", err)
 				return
 			}
 
-			driver := sqlite.NewDriver(db.DBInstance)
 			store := store.New(driver, profile)
 			if err := ExecuteSetup(ctx, store, hostUsername, hostPassword); err != nil {
 				fmt.Printf("failed to setup, error: %+v\n", err)

store/driver.go

@@ -7,9 +7,11 @@ import (
 )
 
 type Driver interface {
+	Close() error
+
+	Migrate(ctx context.Context) error
 	Vacuum(ctx context.Context) error
 	BackupTo(ctx context.Context, filename string) error
-	Close() error
 
 	CreateActivity(ctx context.Context, create *Activity) (*Activity, error)
 

store/db/db.go → store/sqlite/migrate.go

-package db
+package sqlite
 
 import (
 	"context"
-	"database/sql"
 	"embed"
 	"fmt"
 	"io/fs"
@@ -13,7 +12,6 @@ import (
 
 	"github.com/pkg/errors"
 
-	"github.com/usememos/memos/server/profile"
 	"github.com/usememos/memos/server/version"
 )
 
@@ -23,59 +21,14 @@ var migrationFS embed.FS
 //go:embed seed
 var seedFS embed.FS
 
-type DB struct {
-	// sqlite db connection instance
-	DBInstance *sql.DB
-	profile    *profile.Profile
-}
-
-// NewDB returns a new instance of DB associated with the given datasource name.
-func NewDB(profile *profile.Profile) *DB {
-	db := &DB{
-		profile: profile,
-	}
-	return db
-}
-
-// Open opens a database specified by its database driver name and a
-// driver-specific data source name, usually consisting of at least a
-// database name and connection information.
-func (db *DB) Open() error {
-	// Ensure a DSN is set before attempting to open the database.
-	if db.profile.DSN == "" {
-		return errors.New("dsn required")
-	}
-
-	// Connect to the database with some sane settings:
-	// - No shared-cache: it's obsolete; WAL journal mode is a better solution.
-	// - No foreign key constraints: it's currently disabled by default, but it's a
-	// good practice to be explicit and prevent future surprises on SQLite upgrades.
-	// - Journal mode set to WAL: it's the recommended journal mode for most applications
-	// as it prevents locking issues.
-	//
-	// Notes:
-	// - When using the `modernc.org/sqlite` driver, each pragma must be prefixed with `_pragma=`.
-	//
-	// References:
-	// - https://pkg.go.dev/modernc.org/sqlite#Driver.Open
-	// - https://www.sqlite.org/sharedcache.html
-	// - https://www.sqlite.org/pragma.html
-	sqliteDB, err := sql.Open("sqlite", db.profile.DSN+"?_pragma=foreign_keys(0)&_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)")
-	if err != nil {
-		return errors.Wrapf(err, "failed to open db with dsn: %s", db.profile.DSN)
-	}
-	db.DBInstance = sqliteDB
-	return nil
-}
-
 // Migrate applies the latest schema to the database.
-func (db *DB) Migrate(ctx context.Context) error {
-	if db.profile.Mode == "prod" {
-		_, err := os.Stat(db.profile.DSN)
+func (d *Driver) Migrate(ctx context.Context) error {
+	if d.profile.Mode == "prod" {
+		_, err := os.Stat(d.profile.DSN)
 		if err != nil {
 			// If db file not exists, we should create a new one with latest schema.
 			if errors.Is(err, os.ErrNotExist) {
-				if err := db.applyLatestSchema(ctx); err != nil {
+				if err := d.applyLatestSchema(ctx); err != nil {
 					return errors.Wrap(err, "failed to apply latest schema")
 				}
 			} else {
@@ -83,13 +36,13 @@ func (db *DB) Migrate(ctx context.Context) error {
 			}
 		} else {
 			// If db file exists, we should check if we need to migrate the database.
-			currentVersion := version.GetCurrentVersion(db.profile.Mode)
-			migrationHistoryList, err := db.FindMigrationHistoryList(ctx, &MigrationHistoryFind{})
+			currentVersion := version.GetCurrentVersion(d.profile.Mode)
+			migrationHistoryList, err := d.FindMigrationHistoryList(ctx, &MigrationHistoryFind{})
 			if err != nil {
 				return errors.Wrap(err, "failed to find migration history")
 			}
 			if len(migrationHistoryList) == 0 {
-				_, err := db.UpsertMigrationHistory(ctx, &MigrationHistoryUpsert{
+				_, err := d.UpsertMigrationHistory(ctx, &MigrationHistoryUpsert{
 					Version: currentVersion,
 				})
 				if err != nil {
@@ -109,11 +62,11 @@ func (db *DB) Migrate(ctx context.Context) error {
 				minorVersionList := getMinorVersionList()
 
 				// backup the raw database file before migration
-				rawBytes, err := os.ReadFile(db.profile.DSN)
+				rawBytes, err := os.ReadFile(d.profile.DSN)
 				if err != nil {
 					return errors.Wrap(err, "failed to read raw database file")
 				}
-				backupDBFilePath := fmt.Sprintf("%s/memos_%s_%d_backup.db", db.profile.Data, db.profile.Version, time.Now().Unix())
+				backupDBFilePath := fmt.Sprintf("%s/memos_%s_%d_backup.db", d.profile.Data, d.profile.Version, time.Now().Unix())
 				if err := os.WriteFile(backupDBFilePath, rawBytes, 0644); err != nil {
 					return errors.Wrap(err, "failed to write raw database file")
 				}
@@ -124,7 +77,7 @@ func (db *DB) Migrate(ctx context.Context) error {
 					normalizedVersion := minorVersion + ".0"
 					if version.IsVersionGreaterThan(normalizedVersion, latestMigrationHistoryVersion) && version.IsVersionGreaterOrEqualThan(currentVersion, normalizedVersion) {
 						println("applying migration for", normalizedVersion)
-						if err := db.applyMigrationForMinorVersion(ctx, minorVersion); err != nil {
+						if err := d.applyMigrationForMinorVersion(ctx, minorVersion); err != nil {
 							return errors.Wrap(err, "failed to apply minor version migration")
 						}
 					}
@@ -139,13 +92,13 @@ func (db *DB) Migrate(ctx context.Context) error {
 		}
 	} else {
 		// In non-prod mode, we should always migrate the database.
-		if _, err := os.Stat(db.profile.DSN); errors.Is(err, os.ErrNotExist) {
-			if err := db.applyLatestSchema(ctx); err != nil {
+		if _, err := os.Stat(d.profile.DSN); errors.Is(err, os.ErrNotExist) {
+			if err := d.applyLatestSchema(ctx); err != nil {
 				return errors.Wrap(err, "failed to apply latest schema")
 			}
 			// In demo mode, we should seed the database.
-			if db.profile.Mode == "demo" {
-				if err := db.seed(ctx); err != nil {
+			if d.profile.Mode == "demo" {
+				if err := d.seed(ctx); err != nil {
 					return errors.Wrap(err, "failed to seed")
 				}
 			}
@@ -159,9 +112,9 @@ const (
 	latestSchemaFileName = "LATEST__SCHEMA.sql"
 )
 
-func (db *DB) applyLatestSchema(ctx context.Context) error {
+func (d *Driver) applyLatestSchema(ctx context.Context) error {
 	schemaMode := "dev"
-	if db.profile.Mode == "prod" {
+	if d.profile.Mode == "prod" {
 		schemaMode = "prod"
 	}
 	latestSchemaPath := fmt.Sprintf("%s/%s/%s", "migration", schemaMode, latestSchemaFileName)
@@ -170,13 +123,13 @@ func (db *DB) applyLatestSchema(ctx context.Context) error {
 		return errors.Wrapf(err, "failed to read latest schema %q", latestSchemaPath)
 	}
 	stmt := string(buf)
-	if err := db.execute(ctx, stmt); err != nil {
+	if err := d.execute(ctx, stmt); err != nil {
 		return errors.Wrapf(err, "migrate error: %s", stmt)
 	}
 	return nil
 }
 
-func (db *DB) applyMigrationForMinorVersion(ctx context.Context, minorVersion string) error {
+func (d *Driver) applyMigrationForMinorVersion(ctx context.Context, minorVersion string) error {
 	filenames, err := fs.Glob(migrationFS, fmt.Sprintf("%s/%s/*.sql", "migration/prod", minorVersion))
 	if err != nil {
 		return errors.Wrap(err, "failed to read ddl files")
@@ -193,14 +146,14 @@ func (db *DB) applyMigrationForMinorVersion(ctx context.Context, minorVersion st
 		}
 		stmt := string(buf)
 		migrationStmt += stmt
-		if err := db.execute(ctx, stmt); err != nil {
+		if err := d.execute(ctx, stmt); err != nil {
 			return errors.Wrapf(err, "migrate error: %s", stmt)
 		}
 	}
 
 	// Upsert the newest version to migration_history.
 	version := minorVersion + ".0"
-	if _, err = db.UpsertMigrationHistory(ctx, &MigrationHistoryUpsert{
+	if _, err = d.UpsertMigrationHistory(ctx, &MigrationHistoryUpsert{
 		Version: version,
 	}); err != nil {
 		return errors.Wrapf(err, "failed to upsert migration history with version: %s", version)
@@ -209,7 +162,7 @@ func (db *DB) applyMigrationForMinorVersion(ctx context.Context, minorVersion st
 	return nil
 }
 
-func (db *DB) seed(ctx context.Context) error {
+func (d *Driver) seed(ctx context.Context) error {
 	filenames, err := fs.Glob(seedFS, fmt.Sprintf("%s/*.sql", "seed"))
 	if err != nil {
 		return errors.Wrap(err, "failed to read seed files")
@@ -224,7 +177,7 @@ func (db *DB) seed(ctx context.Context) error {
 			return errors.Wrapf(err, "failed to read seed file, filename=%s", filename)
 		}
 		stmt := string(buf)
-		if err := db.execute(ctx, stmt); err != nil {
+		if err := d.execute(ctx, stmt); err != nil {
 			return errors.Wrapf(err, "seed error: %s", stmt)
 		}
 	}
@@ -232,8 +185,8 @@ func (db *DB) seed(ctx context.Context) error {
 }
 
 // execute runs a single SQL statement within a transaction.
-func (db *DB) execute(ctx context.Context, stmt string) error {
-	tx, err := db.DBInstance.Begin()
+func (d *Driver) execute(ctx context.Context, stmt string) error {
+	tx, err := d.db.Begin()
 	if err != nil {
 		return err
 	}

store/db/migration_history.go → store/sqlite/migration_history.go

-package db
+package sqlite
 
 import (
 	"context"
@@ -18,7 +18,7 @@ type MigrationHistoryFind struct {
 	Version *string
 }
 
-func (db *DB) FindMigrationHistoryList(ctx context.Context, find *MigrationHistoryFind) ([]*MigrationHistory, error) {
+func (d *Driver) FindMigrationHistoryList(ctx context.Context, find *MigrationHistoryFind) ([]*MigrationHistory, error) {
 	where, args := []string{"1 = 1"}, []any{}
 
 	if v := find.Version; v != nil {
@@ -34,7 +34,7 @@ func (db *DB) FindMigrationHistoryList(ctx context.Context, find *MigrationHisto
 		WHERE ` + strings.Join(where, " AND ") + `
 		ORDER BY created_ts DESC
 	`
-	rows, err := db.DBInstance.QueryContext(ctx, query, args...)
+	rows, err := d.db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -60,7 +60,7 @@ func (db *DB) FindMigrationHistoryList(ctx context.Context, find *MigrationHisto
 	return list, nil
 }
 
-func (db *DB) UpsertMigrationHistory(ctx context.Context, upsert *MigrationHistoryUpsert) (*MigrationHistory, error) {
+func (d *Driver) UpsertMigrationHistory(ctx context.Context, upsert *MigrationHistoryUpsert) (*MigrationHistory, error) {
 	stmt := `
 		INSERT INTO migration_history (
 			version
@@ -72,7 +72,7 @@ func (db *DB) UpsertMigrationHistory(ctx context.Context, upsert *MigrationHisto
 		RETURNING version, created_ts
 	`
 	var migrationHistory MigrationHistory
-	if err := db.DBInstance.QueryRowContext(ctx, stmt, upsert.Version).Scan(
+	if err := d.db.QueryRowContext(ctx, stmt, upsert.Version).Scan(
 		&migrationHistory.Version,
 		&migrationHistory.CreatedTs,
 	); err != nil {

store/sqlite/sqlite.go

@@ -7,17 +7,46 @@ import (
 	"github.com/pkg/errors"
 	"modernc.org/sqlite"
 
+	"github.com/usememos/memos/server/profile"
 	"github.com/usememos/memos/store"
 )
 
 type Driver struct {
-	db *sql.DB
+	db      *sql.DB
+	profile *profile.Profile
 }
 
-func NewDriver(db *sql.DB) store.Driver {
-	return &Driver{
-		db: db,
+// NewDriver opens a database specified by its database driver name and a
+// driver-specific data source name, usually consisting of at least a
+// database name and connection information.
+func NewDriver(profile *profile.Profile) (store.Driver, error) {
+	// Ensure a DSN is set before attempting to open the database.
+	if profile.DSN == "" {
+		return nil, errors.New("dsn required")
 	}
+
+	// Connect to the database with some sane settings:
+	// - No shared-cache: it's obsolete; WAL journal mode is a better solution.
+	// - No foreign key constraints: it's currently disabled by default, but it's a
+	// good practice to be explicit and prevent future surprises on SQLite upgrades.
+	// - Journal mode set to WAL: it's the recommended journal mode for most applications
+	// as it prevents locking issues.
+	//
+	// Notes:
+	// - When using the `modernc.org/sqlite` driver, each pragma must be prefixed with `_pragma=`.
+	//
+	// References:
+	// - https://pkg.go.dev/modernc.org/sqlite#Driver.Open
+	// - https://www.sqlite.org/sharedcache.html
+	// - https://www.sqlite.org/pragma.html
+	sqliteDB, err := sql.Open("sqlite", profile.DSN+"?_pragma=foreign_keys(0)&_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)")
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to open db with dsn: %s", profile.DSN)
+	}
+
+	driver := Driver{db: sqliteDB, profile: profile}
+
+	return &driver, nil
 }
 
 func (d *Driver) Vacuum(ctx context.Context) error {

test/server/server.go

@@ -18,7 +18,6 @@ import (
 	"github.com/usememos/memos/server"
 	"github.com/usememos/memos/server/profile"
 	"github.com/usememos/memos/store"
-	"github.com/usememos/memos/store/db"
 	"github.com/usememos/memos/store/sqlite"
 	"github.com/usememos/memos/test"
 )
@@ -32,15 +31,14 @@ type TestingServer struct {
 
 func NewTestingServer(ctx context.Context, t *testing.T) (*TestingServer, error) {
 	profile := test.GetTestingProfile(t)
-	db := db.NewDB(profile)
-	if err := db.Open(); err != nil {
-		return nil, errors.Wrap(err, "failed to open db")
+	driver, err := sqlite.NewDriver(profile)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to create db driver")
 	}
-	if err := db.Migrate(ctx); err != nil {
+	if err := driver.Migrate(ctx); err != nil {
 		return nil, errors.Wrap(err, "failed to migrate db")
 	}
 
-	driver := sqlite.NewDriver(db.DBInstance)
 	store := store.New(driver, profile)
 	server, err := server.NewServer(ctx, profile, store)
 	if err != nil {

test/store/store.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	"github.com/usememos/memos/store"
-	"github.com/usememos/memos/store/db"
 	"github.com/usememos/memos/store/sqlite"
 	"github.com/usememos/memos/test"
 
@@ -16,15 +15,14 @@ import (
 
 func NewTestingStore(ctx context.Context, t *testing.T) *store.Store {
 	profile := test.GetTestingProfile(t)
-	db := db.NewDB(profile)
-	if err := db.Open(); err != nil {
-		fmt.Printf("failed to open db, error: %+v\n", err)
+	driver, err := sqlite.NewDriver(profile)
+	if err != nil {
+		fmt.Printf("failed to create db driver, error: %+v\n", err)
 	}
-	if err := db.Migrate(ctx); err != nil {
+	if err := driver.Migrate(ctx); err != nil {
 		fmt.Printf("failed to migrate db, error: %+v\n", err)
 	}
 
-	driver := sqlite.NewDriver(db.DBInstance)
 	store := store.New(driver, profile)
 	return store
 }