PR-Helper/database/db.go

package database

import (
	"database/sql"
	"log"

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

	"github.com/HoHD/PR-Helper/models"
)

type DB struct {
	conn *sql.DB
}

func New(dbPath string) (*DB, error) {
	conn, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_busy_timeout=5000")
	if err != nil {
		return nil, err
	}
	db := &DB{conn: conn}
	if err := db.migrate(); err != nil {
		return nil, err
	}
	if err := db.seedDefaults(); err != nil {
		return nil, err
	}
	return db, nil
}

func (db *DB) Close() error {
	return db.conn.Close()
}

func (db *DB) Conn() *sql.DB {
	return db.conn
}

func (db *DB) migrate() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS settings (
			key   TEXT PRIMARY KEY,
			value TEXT NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS users (
			id            INTEGER PRIMARY KEY AUTOINCREMENT,
			email         TEXT NOT NULL UNIQUE,
			password_hash TEXT NOT NULL,
			created_at    DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at    DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE TABLE IF NOT EXISTS user_settings (
			user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
			key     TEXT NOT NULL,
			value   TEXT NOT NULL,
			PRIMARY KEY (user_id, key)
		)`,
		`CREATE TABLE IF NOT EXISTS repositories (
			id         INTEGER PRIMARY KEY AUTOINCREMENT,
			user_id    INTEGER REFERENCES users(id) ON DELETE CASCADE,
			url        TEXT NOT NULL,
			local_path TEXT NOT NULL,
			size_bytes INTEGER DEFAULT 0,
			cloned_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
			last_used  DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE TABLE IF NOT EXISTS analyses (
			id          INTEGER PRIMARY KEY AUTOINCREMENT,
			user_id     INTEGER REFERENCES users(id) ON DELETE CASCADE,
			repo_id     INTEGER REFERENCES repositories(id),
			type        TEXT NOT NULL,
			base_ref    TEXT NOT NULL,
			head_ref    TEXT NOT NULL,
			result      TEXT,
			created_at  DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE TABLE IF NOT EXISTS review_notes (
			id          INTEGER PRIMARY KEY AUTOINCREMENT,
			user_id     INTEGER REFERENCES users(id) ON DELETE CASCADE,
			analysis_id INTEGER REFERENCES analyses(id),
			scope       TEXT NOT NULL,
			scope_key   TEXT NOT NULL,
			content     TEXT NOT NULL,
			created_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at  DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
	}
	for _, s := range stmts {
		if _, err := db.conn.Exec(s); err != nil {
			return err
		}
	}
	// Run incremental migrations for existing databases
	if err := db.migrateV2(); err != nil {
		return err
	}
	return nil
}

// migrateV2 adds user_id columns to existing tables and removes the old UNIQUE constraint on url.
func (db *DB) migrateV2() error {
	migrations := []string{
		`ALTER TABLE repositories ADD COLUMN user_id INTEGER REFERENCES users(id) ON DELETE CASCADE`,
		`ALTER TABLE analyses ADD COLUMN user_id INTEGER REFERENCES users(id) ON DELETE CASCADE`,
		`ALTER TABLE review_notes ADD COLUMN user_id INTEGER REFERENCES users(id) ON DELETE CASCADE`,
	}
	for _, m := range migrations {
		// ALTER TABLE ADD COLUMN will fail if column already exists; ignore the error
		db.conn.Exec(m)
	}
	// Recreate repositories table without the UNIQUE constraint on url (same url can be cloned by different users)
	// SQLite doesn't support DROP CONSTRAINT, so we use a workaround
	var hasOldUnique int
	db.conn.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='repositories_old'`).Scan(&hasOldUnique)
	if hasOldUnique == 0 {
		// Check if the UNIQUE constraint exists by trying to detect the old schema
		// We'll just do the migration idempotently
		db.conn.Exec(`CREATE TABLE IF NOT EXISTS repositories_new (
			id         INTEGER PRIMARY KEY AUTOINCREMENT,
			user_id    INTEGER REFERENCES users(id) ON DELETE CASCADE,
			url        TEXT NOT NULL,
			local_path TEXT NOT NULL,
			size_bytes INTEGER DEFAULT 0,
			cloned_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
			last_used  DATETIME DEFAULT CURRENT_TIMESTAMP
		)`)
		db.conn.Exec(`INSERT OR IGNORE INTO repositories_new (id, user_id, url, local_path, size_bytes, cloned_at, last_used)
			SELECT id, user_id, url, local_path, size_bytes, cloned_at, last_used FROM repositories`)
		// Only swap if the new table has data or old table is empty
		var newCount, oldCount int
		db.conn.QueryRow(`SELECT COUNT(*) FROM repositories_new`).Scan(&newCount)
		db.conn.QueryRow(`SELECT COUNT(*) FROM repositories`).Scan(&oldCount)
		if newCount >= oldCount {
			db.conn.Exec(`DROP TABLE repositories`)
			db.conn.Exec(`ALTER TABLE repositories_new RENAME TO repositories`)
		} else {
			db.conn.Exec(`DROP TABLE repositories_new`)
		}
	}
	return nil
}

func (db *DB) seedDefaults() error {
	for key, val := range models.DefaultSettings {
		_, err := db.conn.Exec(
			`INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)`, key, val,
		)
		if err != nil {
			log.Printf("warning: failed to seed setting %s: %v", key, err)
		}
	}
	return nil
}