Files
MeBox/internal/database/database.go
T
2026-06-07 18:30:27 +08:00

123 lines
3.2 KiB
Go

// Package database wires up GORM against SQLite (WAL mode) and exposes the
// auto-migration entry point used at startup.
package database
import (
"fmt"
"path/filepath"
"github.com/glebarez/sqlite"
"go.uber.org/zap"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"github.com/ShukeBta/MediaStationGo/internal/config"
"github.com/ShukeBta/MediaStationGo/internal/model"
)
// Open initialises the SQLite database file applying WAL pragmas for
// better concurrent read performance — same defaults as nowen-video.
func Open(cfg *config.Config, log *zap.Logger) (*gorm.DB, error) {
dsn := buildDSN(cfg)
gormLogger := logger.New(
zapStdLogger{log: log},
logger.Config{
SlowThreshold: 0,
LogLevel: logger.Warn,
IgnoreRecordNotFoundError: true,
Colorful: false,
},
)
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{
Logger: gormLogger,
PrepareStmt: true,
DisableForeignKeyConstraintWhenMigrating: false,
})
if err != nil {
return nil, fmt.Errorf("gorm open: %w", err)
}
sqlDB, err := db.DB()
if err != nil {
return nil, fmt.Errorf("gorm sqldb: %w", err)
}
if cfg.Database.MaxOpenConns > 0 {
sqlDB.SetMaxOpenConns(cfg.Database.MaxOpenConns)
}
if cfg.Database.MaxIdleConns > 0 {
sqlDB.SetMaxIdleConns(cfg.Database.MaxIdleConns)
}
return db, nil
}
func buildDSN(cfg *config.Config) string {
dbPath := cfg.Database.DBPath
if !filepath.IsAbs(dbPath) {
// keep as-is to respect user-provided relative paths.
dbPath = filepath.Clean(dbPath)
}
dsn := dbPath + "?_pragma=foreign_keys(1)"
if cfg.Database.WALMode {
dsn += "&_pragma=journal_mode(WAL)"
}
if cfg.Database.BusyTimeout > 0 {
dsn += fmt.Sprintf("&_pragma=busy_timeout(%d)", cfg.Database.BusyTimeout)
}
if cfg.Database.CacheSize != 0 {
dsn += fmt.Sprintf("&_pragma=cache_size(%d)", cfg.Database.CacheSize)
}
return dsn
}
// AutoMigrate creates tables for every model registered in the model package.
func AutoMigrate(db *gorm.DB) error {
if err := db.AutoMigrate(model.AllModels()...); err != nil {
return err
}
return enforceTelegramBindingOneToOne(db)
}
func enforceTelegramBindingOneToOne(db *gorm.DB) error {
if !db.Migrator().HasTable(&model.TelegramBinding{}) {
return nil
}
return db.Transaction(func(tx *gorm.DB) error {
if err := tx.Exec(`
DELETE FROM telegram_bindings
WHERE deleted_at IS NULL
AND user_id IN (
SELECT user_id
FROM telegram_bindings
WHERE deleted_at IS NULL
GROUP BY user_id
HAVING COUNT(*) > 1
)
AND id NOT IN (
SELECT id
FROM (
SELECT id,
ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY created_at ASC, id ASC) AS rn
FROM telegram_bindings
WHERE deleted_at IS NULL
)
WHERE rn = 1
)
`).Error; err != nil {
return err
}
return tx.Exec(`
CREATE UNIQUE INDEX IF NOT EXISTS idx_telegram_bindings_user_id_active
ON telegram_bindings(user_id)
WHERE deleted_at IS NULL
`).Error
})
}
// zapStdLogger adapts a *zap.Logger to GORM's tiny logger interface.
type zapStdLogger struct{ log *zap.Logger }
func (z zapStdLogger) Printf(format string, args ...interface{}) {
z.log.Sugar().Infof(format, args...)
}