streamrip-go/internal/store/sqlite.go

package store

import (
	"context"
	"database/sql"
	"fmt"
	"io"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

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

	// hasLegacyRows is set when the database was migrated from the pre-source
	// schema. Those rows have no source, so IsDownloaded has to consider them
	// separately or the entire pre-migration library looks undownloaded.
	hasLegacyRows bool
	legacyWarned  bool
}

func NewSQLite(path string) (*SQLite, error) {
	db, err := sql.Open("sqlite", sqliteDSN(path))
	if err != nil {
		return nil, err
	}
	// One writer at a time. SQLite allows a single writer anyway, and the
	// mutex below already serializes this process; capping the pool keeps the
	// driver from opening connections that can only contend with each other.
	db.SetMaxOpenConns(1)

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

// sqliteDSN builds a file: URI with the pragmas this app needs.
//
// WAL lets a concurrent reader (for example `rip db browse`) run while a rip
// is writing, and busy_timeout makes a contended write wait instead of failing
// immediately with SQLITE_BUSY. The path is passed through url.URL so that a
// config path containing '?' or '#' is not parsed as query parameters.
func sqliteDSN(path string) string {
	u := url.URL{Scheme: "file", Opaque: (&url.URL{Path: path}).EscapedPath()}
	q := url.Values{}
	q.Add("_pragma", "busy_timeout(10000)")
	q.Add("_pragma", "journal_mode(WAL)")
	q.Add("_pragma", "foreign_keys(on)")
	return u.String() + "?" + q.Encode()
}

// schemaVersion is stored in PRAGMA user_version so later migrations can be
// ordered explicitly instead of being inferred from PRAGMA table_info.
const schemaVersion = 1

func (s *SQLite) userVersion() (int, error) {
	var v int
	if err := s.db.QueryRow(`PRAGMA user_version`).Scan(&v); err != nil {
		return 0, err
	}
	return v, nil
}

func (s *SQLite) setUserVersion(v int) error {
	_, err := s.db.Exec(fmt.Sprintf(`PRAGMA user_version = %d`, v))
	return err
}

// backupFile copies the database aside before a destructive migration. A
// migration that loses the download history silently costs the user a full
// re-download, so keep a copy they can fall back to.
func (s *SQLite) backupFile() (string, error) {
	if strings.TrimSpace(s.path) == "" {
		return "", nil
	}
	src, err := os.Open(s.path)
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", err
	}
	defer func() { _ = src.Close() }()

	dstPath := fmt.Sprintf("%s.bak-%d", s.path, time.Now().Unix())
	dst, err := os.OpenFile(dstPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return "", err
	}
	if _, err = io.Copy(dst, src); err != nil {
		_ = dst.Close()
		_ = os.Remove(dstPath)
		return "", err
	}
	if err = dst.Close(); err != nil {
		return "", err
	}
	return dstPath, 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
		}
	}

	// A database migrated by an earlier version has sourceless rows but no
	// in-memory flag, so detect them here rather than only at migration time.
	if !s.hasLegacyRows {
		if n, err := s.legacyRowCount(context.Background()); err == nil && n > 0 {
			s.hasLegacyRows = true
		}
	}

	if v, err := s.userVersion(); err == nil && v < schemaVersion {
		if err = s.setUserVersion(schemaVersion); 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
	}

	// Legacy schema: a downloads table with no source column. Back the file up
	// first, because the rows below can only be carried over without a source
	// and that is not fully reversible.
	backupPath, backupErr := s.backupFile()
	if backupErr != nil {
		return fmt.Errorf("could not back up %s before migrating: %w", s.path, backupErr)
	}

	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
	}
	// downloads_legacy is deliberately kept. Dropping it used to make the
	// migration unrecoverable, and IsDownloaded needs the sourceless rows to
	// avoid re-downloading everything the user already has.

	if err = tx.Commit(); err != nil {
		return err
	}

	s.hasLegacyRows = true
	if backupPath != "" {
		fmt.Fprintf(os.Stderr, "download database migrated to the source-aware schema; previous copy kept at %s\n", backupPath)
	}
	return nil
}

// legacyRowCount reports how many pre-migration rows exist. Rows carried over
// from the old schema have an empty source.
func (s *SQLite) legacyRowCount(ctx context.Context) (int, error) {
	var n int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(1) FROM downloads WHERE source = ''`).Scan(&n)
	return n, 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
	}
	if count > 0 {
		return true, nil
	}

	// Rows migrated from the pre-source schema have an empty source, so an
	// exact match on (source, id) misses them and the whole pre-migration
	// library would look undownloaded. Matching on id alone can only ever
	// over-match across sources, which costs a skipped re-download rather than
	// a duplicated library.
	if !s.hasLegacyRows {
		return false, nil
	}
	err = s.db.QueryRowContext(ctx, `SELECT COUNT(1) FROM downloads WHERE source = '' AND id = ?`, id).Scan(&count)
	if err != nil {
		return false, err
	}
	if count > 0 && !s.legacyWarned {
		s.legacyWarned = true
		fmt.Fprintf(os.Stderr, "note: matching downloads recorded before sources were tracked; pass --force to re-download them\n")
	}
	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()
}