jiggablend/internal/database/schema.go

package database

import (
	"database/sql"
	"fmt"

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

// DB wraps the database connection
type DB struct {
	*sql.DB
}

// NewDB creates a new database connection
func NewDB(dbPath string) (*DB, error) {
	db, err := sql.Open("sqlite3", dbPath+"?_foreign_keys=1")
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	if err := db.Ping(); err != nil {
		return nil, fmt.Errorf("failed to ping database: %w", err)
	}

	database := &DB{DB: db}
	if err := database.migrate(); err != nil {
		return nil, fmt.Errorf("failed to migrate database: %w", err)
	}

	return database, nil
}

// migrate runs database migrations
func (db *DB) migrate() error {
	schema := `
	CREATE TABLE IF NOT EXISTS users (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		email TEXT UNIQUE NOT NULL,
		name TEXT NOT NULL,
		oauth_provider TEXT NOT NULL,
		oauth_id TEXT NOT NULL,
		is_admin BOOLEAN NOT NULL DEFAULT 0,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		UNIQUE(oauth_provider, oauth_id)
	);

	CREATE TABLE IF NOT EXISTS jobs (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		user_id INTEGER NOT NULL,
		name TEXT NOT NULL,
		status TEXT NOT NULL DEFAULT 'pending',
		progress REAL NOT NULL DEFAULT 0.0,
		frame_start INTEGER NOT NULL,
		frame_end INTEGER NOT NULL,
		output_format TEXT NOT NULL DEFAULT 'PNG',
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		started_at DATETIME,
		completed_at DATETIME,
		error_message TEXT,
		FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
	);

	CREATE TABLE IF NOT EXISTS runners (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		name TEXT NOT NULL,
		hostname TEXT NOT NULL,
		ip_address TEXT NOT NULL,
		status TEXT NOT NULL DEFAULT 'offline',
		last_heartbeat DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		capabilities TEXT,
		registration_token TEXT,
		runner_secret TEXT,
		manager_secret TEXT,
		verified BOOLEAN NOT NULL DEFAULT 0,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
	);

	CREATE TABLE IF NOT EXISTS tasks (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		job_id INTEGER NOT NULL,
		runner_id INTEGER,
		frame_start INTEGER NOT NULL,
		frame_end INTEGER NOT NULL,
		status TEXT NOT NULL DEFAULT 'pending',
		output_path TEXT,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		started_at DATETIME,
		completed_at DATETIME,
		error_message TEXT,
		FOREIGN KEY (job_id) REFERENCES jobs(id) ON DELETE CASCADE,
		FOREIGN KEY (runner_id) REFERENCES runners(id) ON DELETE SET NULL
	);

	CREATE TABLE IF NOT EXISTS job_files (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		job_id INTEGER NOT NULL,
		file_type TEXT NOT NULL,
		file_path TEXT NOT NULL,
		file_name TEXT NOT NULL,
		file_size INTEGER NOT NULL,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		FOREIGN KEY (job_id) REFERENCES jobs(id) ON DELETE CASCADE
	);

	CREATE TABLE IF NOT EXISTS manager_secrets (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		secret TEXT UNIQUE NOT NULL,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
	);

	CREATE TABLE IF NOT EXISTS registration_tokens (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		token TEXT UNIQUE NOT NULL,
		expires_at DATETIME NOT NULL,
		used BOOLEAN NOT NULL DEFAULT 0,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		created_by INTEGER,
		FOREIGN KEY (created_by) REFERENCES users(id) ON DELETE SET NULL
	);

	CREATE TABLE IF NOT EXISTS task_logs (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		task_id INTEGER NOT NULL,
		runner_id INTEGER,
		log_level TEXT NOT NULL,
		message TEXT NOT NULL,
		step_name TEXT,
		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE,
		FOREIGN KEY (runner_id) REFERENCES runners(id) ON DELETE SET NULL
	);

	CREATE TABLE IF NOT EXISTS task_steps (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		task_id INTEGER NOT NULL,
		step_name TEXT NOT NULL,
		status TEXT NOT NULL DEFAULT 'pending',
		started_at DATETIME,
		completed_at DATETIME,
		duration_ms INTEGER,
		error_message TEXT,
		FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE
	);

	CREATE INDEX IF NOT EXISTS idx_jobs_user_id ON jobs(user_id);
	CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status);
	CREATE INDEX IF NOT EXISTS idx_tasks_job_id ON tasks(job_id);
	CREATE INDEX IF NOT EXISTS idx_tasks_runner_id ON tasks(runner_id);
	CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
	CREATE INDEX IF NOT EXISTS idx_tasks_started_at ON tasks(started_at);
	CREATE INDEX IF NOT EXISTS idx_job_files_job_id ON job_files(job_id);
	CREATE INDEX IF NOT EXISTS idx_registration_tokens_token ON registration_tokens(token);
	CREATE INDEX IF NOT EXISTS idx_registration_tokens_expires_at ON registration_tokens(expires_at);
	CREATE INDEX IF NOT EXISTS idx_task_logs_task_id_created_at ON task_logs(task_id, created_at);
	CREATE INDEX IF NOT EXISTS idx_task_logs_runner_id ON task_logs(runner_id);
	CREATE INDEX IF NOT EXISTS idx_task_steps_task_id ON task_steps(task_id);
	CREATE INDEX IF NOT EXISTS idx_runners_last_heartbeat ON runners(last_heartbeat);
	`

	if _, err := db.Exec(schema); err != nil {
		return fmt.Errorf("failed to create schema: %w", err)
	}

	// Migrate existing tables to add new columns
	migrations := []string{
		// Add is_admin to users if it doesn't exist
		`ALTER TABLE users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT 0`,
		// Add new columns to runners if they don't exist
		`ALTER TABLE runners ADD COLUMN registration_token TEXT`,
		`ALTER TABLE runners ADD COLUMN runner_secret TEXT`,
		`ALTER TABLE runners ADD COLUMN manager_secret TEXT`,
		`ALTER TABLE runners ADD COLUMN verified BOOLEAN NOT NULL DEFAULT 0`,
		// Add allow_parallel_runners to jobs if it doesn't exist
		`ALTER TABLE jobs ADD COLUMN allow_parallel_runners BOOLEAN NOT NULL DEFAULT 1`,
		// Add timeout_seconds to jobs if it doesn't exist
		`ALTER TABLE jobs ADD COLUMN timeout_seconds INTEGER DEFAULT 86400`,
		// Add blend_metadata to jobs if it doesn't exist
		`ALTER TABLE jobs ADD COLUMN blend_metadata TEXT`,
		// Add task_type to tasks if it doesn't exist
		`ALTER TABLE tasks ADD COLUMN task_type TEXT DEFAULT 'render'`,
		// Add new columns to tasks if they don't exist
		`ALTER TABLE tasks ADD COLUMN current_step TEXT`,
		`ALTER TABLE tasks ADD COLUMN retry_count INTEGER DEFAULT 0`,
		`ALTER TABLE tasks ADD COLUMN max_retries INTEGER DEFAULT 3`,
		`ALTER TABLE tasks ADD COLUMN timeout_seconds INTEGER`,
	}

	for _, migration := range migrations {
		// SQLite doesn't support IF NOT EXISTS for ALTER TABLE, so we ignore errors
		db.Exec(migration)
	}

	return nil
}

// Close closes the database connection
func (db *DB) Close() error {
	return db.DB.Close()
}