streamrip-go/internal/store/sqlite.go

package store

import (
	"context"
	"database/sql"
	"sync"

	_ "modernc.org/sqlite"
)

type SQLite struct {
	db *sql.DB
	mu sync.Mutex
}

func NewSQLite(path string) (*SQLite, error) {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, err
	}

	s := &SQLite{db: db}
	if err = s.init(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return s, nil
}

func (s *SQLite) init() error {
	if err := s.ensureDownloadsSchema(); err != nil {
		return err
	}

	queries := []string{
		`CREATE TABLE IF NOT EXISTS failed_downloads (
			source TEXT NOT NULL,
			media_type TEXT NOT NULL,
			id TEXT NOT NULL,
			PRIMARY KEY (source, media_type, id)
		)`,
	}

	for _, q := range queries {
		if _, err := s.db.Exec(q); err != nil {
			return err
		}
	}

	return nil
}

func (s *SQLite) ensureDownloadsSchema() error {
	rows, err := s.db.Query(`PRAGMA table_info(downloads)`)
	if err != nil {
		return err
	}
	defer rows.Close()

	hasTable := false
	hasSource := false
	for rows.Next() {
		hasTable = true
		var (
			cid       int
			name      string
			colType   string
			notNull   int
			defaultV  sql.NullString
			pkOrdinal int
		)
		if err = rows.Scan(&cid, &name, &colType, &notNull, &defaultV, &pkOrdinal); err != nil {
			return err
		}
		if name == "source" {
			hasSource = true
		}
	}
	if err = rows.Err(); err != nil {
		return err
	}

	if !hasTable {
		_, err = s.db.Exec(`CREATE TABLE downloads (
			source TEXT NOT NULL,
			id TEXT NOT NULL,
			PRIMARY KEY (source, id)
		)`)
		return err
	}

	if hasSource {
		_, err = s.db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_downloads_source_id ON downloads(source, id)`)
		return err
	}

	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			_ = tx.Rollback()
		}
	}()

	if _, err = tx.Exec(`DROP TABLE IF EXISTS downloads_legacy`); err != nil {
		return err
	}
	if _, err = tx.Exec(`ALTER TABLE downloads RENAME TO downloads_legacy`); err != nil {
		return err
	}
	if _, err = tx.Exec(`CREATE TABLE downloads (
		source TEXT NOT NULL,
		id TEXT NOT NULL,
		PRIMARY KEY (source, id)
	)`); err != nil {
		return err
	}
	if _, err = tx.Exec(`INSERT OR IGNORE INTO downloads(source, id) SELECT '', id FROM downloads_legacy`); err != nil {
		return err
	}
	if _, err = tx.Exec(`DROP TABLE downloads_legacy`); err != nil {
		return err
	}

	err = tx.Commit()
	return err
}

func (s *SQLite) IsDownloaded(ctx context.Context, source, id string) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	var count int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(1) FROM downloads WHERE source = ? AND id = ?`, source, id).Scan(&count)
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

func (s *SQLite) MarkDownloaded(ctx context.Context, source, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, err := s.db.ExecContext(ctx, `INSERT OR IGNORE INTO downloads(source, id) VALUES (?, ?)`, source, id)
	return err
}

func (s *SQLite) MarkFailed(ctx context.Context, source, mediaType, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, err := s.db.ExecContext(
		ctx,
		`INSERT OR IGNORE INTO failed_downloads(source, media_type, id) VALUES (?, ?, ?)`,
		source,
		mediaType,
		id,
	)
	return err
}

func (s *SQLite) Close() error {
	if s.db == nil {
		return nil
	}
	return s.db.Close()
}