byop-engine/dbmanager/sqlite.go

package dbmanager

import (
	"fmt"
	"os"

	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

// SQLiteManager implements the DbManager interface for SQLite using GORM
type SQLiteManager struct {
	db  *gorm.DB
	dsn string
}

// NewSQLiteManager initializes a new SQLiteManager
func NewSQLiteManager(dataSourceName string) (*SQLiteManager, error) {
	// First check if the database file exists
	isNewDb := !fileExists(dataSourceName)
	if isNewDb {
		// Create the database file if it doesn't exist
		file, err := os.Create(dataSourceName)
		if err != nil {
			return nil, fmt.Errorf("failed to create SQLite database file: %w", err)
		}
		defer file.Close()
	}

	// Open SQLite database with GORM and SQLite-specific configuration
	db, err := gorm.Open(sqlite.Open(dataSourceName), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
		// Set DisableForeignKeyConstraintWhenMigrating to true to avoid foreign key issues during migration
		DisableForeignKeyConstraintWhenMigrating: true,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to SQLite database: %w", err)
	}

	// Enable foreign keys in SQLite after migrations
	db.Exec("PRAGMA foreign_keys = ON")

	return &SQLiteManager{
		db:  db,
		dsn: dataSourceName,
	}, nil
}

// fileExists checks if a file exists
func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return !os.IsNotExist(err)
}

// GetDB returns the GORM database instance
func (m *SQLiteManager) GetDB() *gorm.DB {
	return m.db
}

// Connect establishes a connection to the SQLite database
func (m *SQLiteManager) Connect() error {
	// Connection is already established in NewSQLiteManager
	return nil
}

// Disconnect closes the connection to the SQLite database
func (m *SQLiteManager) Disconnect() error {
	sqlDB, err := m.db.DB()
	if err != nil {
		return err
	}
	return sqlDB.Close()
}

// Migrate runs auto migration for the provided models
func (m *SQLiteManager) Migrate(models ...interface{}) error {
	// Check if database file exists before
	isNewDb := !fileExists(m.dsn)

	// SQLite has limitations with ALTER TABLE for adding NOT NULL columns
	// For a new database, we can just create tables directly
	if isNewDb {
		return m.db.AutoMigrate(models...)
	}

	// For existing database, we need a more careful approach to avoid NOT NULL errors
	// Try to create tables that don't exist
	migrator := m.db.Migrator()

	for _, model := range models {
		// Check if the table exists
		if !migrator.HasTable(model) {
			// If table doesn't exist, create it
			if err := migrator.CreateTable(model); err != nil {
				return fmt.Errorf("failed to create table for %T: %w", model, err)
			}
		}
	}

	// Re-enable foreign key constraints
	m.db.Exec("PRAGMA foreign_keys = ON")

	return nil
}