package database import ( "database/sql" "fmt" "log" _ "github.com/marcboeker/go-duckdb/v2" ) // 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("duckdb", dbPath) 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 { // Create sequences for auto-incrementing primary keys sequences := []string{ `CREATE SEQUENCE IF NOT EXISTS seq_users_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_jobs_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_runners_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_tasks_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_job_files_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_manager_secrets_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_registration_tokens_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_task_logs_id START 1`, `CREATE SEQUENCE IF NOT EXISTS seq_task_steps_id START 1`, } for _, seq := range sequences { if _, err := db.Exec(seq); err != nil { return fmt.Errorf("failed to create sequence: %w", err) } } schema := ` CREATE TABLE IF NOT EXISTS users ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_users_id'), email TEXT UNIQUE NOT NULL, name TEXT NOT NULL, oauth_provider TEXT NOT NULL, oauth_id TEXT NOT NULL, password_hash TEXT, is_admin BOOLEAN NOT NULL DEFAULT false, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(oauth_provider, oauth_id) ); CREATE TABLE IF NOT EXISTS jobs ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_jobs_id'), user_id BIGINT NOT NULL, job_type TEXT NOT NULL DEFAULT 'render', name TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', progress REAL NOT NULL DEFAULT 0.0, frame_start INTEGER, frame_end INTEGER, output_format TEXT, allow_parallel_runners BOOLEAN, timeout_seconds INTEGER DEFAULT 86400, blend_metadata TEXT, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, started_at TIMESTAMP, completed_at TIMESTAMP, error_message TEXT, FOREIGN KEY (user_id) REFERENCES users(id) ); CREATE TABLE IF NOT EXISTS runners ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_runners_id'), name TEXT NOT NULL, hostname TEXT NOT NULL, ip_address TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'offline', last_heartbeat TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, capabilities TEXT, registration_token TEXT, runner_secret TEXT, manager_secret TEXT, verified BOOLEAN NOT NULL DEFAULT false, priority INTEGER NOT NULL DEFAULT 100, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS tasks ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_tasks_id'), job_id BIGINT NOT NULL, runner_id BIGINT, frame_start INTEGER NOT NULL, frame_end INTEGER NOT NULL, status TEXT NOT NULL DEFAULT 'pending', output_path TEXT, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, started_at TIMESTAMP, completed_at TIMESTAMP, error_message TEXT ); CREATE TABLE IF NOT EXISTS job_files ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_job_files_id'), job_id BIGINT NOT NULL, file_type TEXT NOT NULL, file_path TEXT NOT NULL, file_name TEXT NOT NULL, file_size BIGINT NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS manager_secrets ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_manager_secrets_id'), secret TEXT UNIQUE NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS registration_tokens ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_registration_tokens_id'), token TEXT UNIQUE NOT NULL, expires_at TIMESTAMP NOT NULL, used BOOLEAN NOT NULL DEFAULT false, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, created_by BIGINT, FOREIGN KEY (created_by) REFERENCES users(id) ); CREATE TABLE IF NOT EXISTS task_logs ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_task_logs_id'), task_id BIGINT NOT NULL, runner_id BIGINT, log_level TEXT NOT NULL, message TEXT NOT NULL, step_name TEXT, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS task_steps ( id BIGINT PRIMARY KEY DEFAULT nextval('seq_task_steps_id'), task_id BIGINT NOT NULL, step_name TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', started_at TIMESTAMP, completed_at TIMESTAMP, duration_ms INTEGER, error_message TEXT ); 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); CREATE TABLE IF NOT EXISTS settings ( key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); ` 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 IF NOT EXISTS is_admin BOOLEAN NOT NULL DEFAULT false`, // Add new columns to runners if they don't exist `ALTER TABLE runners ADD COLUMN IF NOT EXISTS registration_token TEXT`, `ALTER TABLE runners ADD COLUMN IF NOT EXISTS runner_secret TEXT`, `ALTER TABLE runners ADD COLUMN IF NOT EXISTS manager_secret TEXT`, `ALTER TABLE runners ADD COLUMN IF NOT EXISTS verified BOOLEAN NOT NULL DEFAULT false`, `ALTER TABLE runners ADD COLUMN IF NOT EXISTS priority INTEGER NOT NULL DEFAULT 100`, // Add allow_parallel_runners to jobs if it doesn't exist `ALTER TABLE jobs ADD COLUMN IF NOT EXISTS allow_parallel_runners BOOLEAN NOT NULL DEFAULT true`, // Add timeout_seconds to jobs if it doesn't exist `ALTER TABLE jobs ADD COLUMN IF NOT EXISTS timeout_seconds INTEGER DEFAULT 86400`, // Add blend_metadata to jobs if it doesn't exist `ALTER TABLE jobs ADD COLUMN IF NOT EXISTS blend_metadata TEXT`, // Add job_type to jobs if it doesn't exist `ALTER TABLE jobs ADD COLUMN IF NOT EXISTS job_type TEXT DEFAULT 'render'`, // Add task_type to tasks if it doesn't exist `ALTER TABLE tasks ADD COLUMN IF NOT EXISTS task_type TEXT DEFAULT 'render'`, // Add new columns to tasks if they don't exist `ALTER TABLE tasks ADD COLUMN IF NOT EXISTS current_step TEXT`, `ALTER TABLE tasks ADD COLUMN IF NOT EXISTS retry_count INTEGER DEFAULT 0`, `ALTER TABLE tasks ADD COLUMN IF NOT EXISTS max_retries INTEGER DEFAULT 3`, `ALTER TABLE tasks ADD COLUMN IF NOT EXISTS timeout_seconds INTEGER`, // Migrate file_size from INTEGER to BIGINT to support large files (>2GB) // DuckDB doesn't support direct ALTER COLUMN TYPE, so we use a workaround: // 1. Add new column as BIGINT // 2. Copy data from old column // 3. Drop old column // 4. Rename new column // Note: This will only run if the column exists and is INTEGER `ALTER TABLE job_files ADD COLUMN IF NOT EXISTS file_size_new BIGINT`, `UPDATE job_files SET file_size_new = CAST(file_size AS BIGINT) WHERE file_size_new IS NULL`, `ALTER TABLE job_files DROP COLUMN IF EXISTS file_size`, `ALTER TABLE job_files RENAME COLUMN file_size_new TO file_size`, } for _, migration := range migrations { // DuckDB supports IF NOT EXISTS for ALTER TABLE, so we can safely execute if _, err := db.Exec(migration); err != nil { // Log but don't fail - column might already exist or table might not exist yet // This is fine for migrations that run after schema creation // For the file_size migration, if it fails (e.g., already BIGINT), that's fine } } // Initialize registration_enabled setting (default: true) if it doesn't exist var settingCount int err := db.QueryRow("SELECT COUNT(*) FROM settings WHERE key = ?", "registration_enabled").Scan(&settingCount) if err == nil && settingCount == 0 { _, err = db.Exec("INSERT INTO settings (key, value) VALUES (?, ?)", "registration_enabled", "true") if err != nil { // Log but don't fail - setting might have been created by another process log.Printf("Note: Could not initialize registration_enabled setting: %v", err) } } return nil } // Close closes the database connection func (db *DB) Close() error { return db.DB.Close() }