stremio/node/media-node/internal/database/store.go

package database

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"time"

	_ "modernc.org/sqlite"
)

type Store struct {
	db *sql.DB
}

type LocalFile struct {
	LocalFileID string
	Path        string
	SizeBytes   int64
	ModifiedAt  time.Time
	MediaType   string
	ReleaseName string
	Revision    int64
}

type PlaybackSession struct {
	SessionID          string
	TokenHash          string
	LocalFileID        string
	IdleTimeoutSeconds int
	AbsoluteExpiresAt  time.Time
	LastActivity       time.Time
}

func Open(dataDir string) (*Store, error) {
	if err := os.MkdirAll(dataDir, 0o750); err != nil {
		return nil, err
	}
	dbPath := filepath.Join(dataDir, "node.db")
	db, err := sql.Open("sqlite", dbPath+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
	if err != nil {
		return nil, err
	}
	s := &Store{db: db}
	if err := s.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return s, nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

func (s *Store) migrate() error {
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS local_files (
			local_file_id TEXT PRIMARY KEY,
			path TEXT NOT NULL UNIQUE,
			size_bytes INTEGER NOT NULL,
			modified_at TEXT NOT NULL,
			media_type TEXT NOT NULL,
			release_name TEXT NOT NULL,
			revision INTEGER NOT NULL DEFAULT 0
		);
		CREATE TABLE IF NOT EXISTS playback_sessions (
			session_id TEXT PRIMARY KEY,
			token_hash TEXT NOT NULL UNIQUE,
			local_file_id TEXT NOT NULL,
			idle_timeout_seconds INTEGER NOT NULL,
			absolute_expires_at TEXT NOT NULL,
			last_activity TEXT NOT NULL
		);
		CREATE TABLE IF NOT EXISTS node_state (
			key TEXT PRIMARY KEY,
			value TEXT NOT NULL
		);
		CREATE INDEX IF NOT EXISTS idx_sessions_token ON playback_sessions(token_hash);
	`)
	return err
}

func (s *Store) UpsertFile(f LocalFile) error {
	_, err := s.db.Exec(`
		INSERT INTO local_files (local_file_id, path, size_bytes, modified_at, media_type, release_name, revision)
		VALUES (?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(local_file_id) DO UPDATE SET
			path=excluded.path,
			size_bytes=excluded.size_bytes,
			modified_at=excluded.modified_at,
			media_type=excluded.media_type,
			release_name=excluded.release_name,
			revision=excluded.revision
	`, f.LocalFileID, f.Path, f.SizeBytes, f.ModifiedAt.UTC().Format(time.RFC3339),
		f.MediaType, f.ReleaseName, f.Revision)
	return err
}

func (s *Store) DeleteFile(localFileID string) error {
	_, err := s.db.Exec(`DELETE FROM local_files WHERE local_file_id = ?`, localFileID)
	return err
}

func (s *Store) GetFile(localFileID string) (*LocalFile, error) {
	row := s.db.QueryRow(`
		SELECT local_file_id, path, size_bytes, modified_at, media_type, release_name, revision
		FROM local_files WHERE local_file_id = ?
	`, localFileID)
	return scanFile(row)
}

func (s *Store) GetFileByPath(path string) (*LocalFile, error) {
	row := s.db.QueryRow(`
		SELECT local_file_id, path, size_bytes, modified_at, media_type, release_name, revision
		FROM local_files WHERE path = ?
	`, path)
	return scanFile(row)
}

func (s *Store) AllFiles() ([]LocalFile, error) {
	rows, err := s.db.Query(`
		SELECT local_file_id, path, size_bytes, modified_at, media_type, release_name, revision
		FROM local_files ORDER BY release_name
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var files []LocalFile
	for rows.Next() {
		f, err := scanFile(rows)
		if err != nil {
			return nil, err
		}
		files = append(files, *f)
	}
	return files, rows.Err()
}

// FilesUnderPath returns the exact path and any files nested under it.
func (s *Store) FilesUnderPath(path string) ([]LocalFile, error) {
	clean := filepath.Clean(path)
	like := clean + string(filepath.Separator) + "%"
	rows, err := s.db.Query(`
		SELECT local_file_id, path, size_bytes, modified_at, media_type, release_name, revision
		FROM local_files
		WHERE path = ? OR path LIKE ?
		ORDER BY path
	`, clean, like)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var files []LocalFile
	for rows.Next() {
		f, err := scanFile(rows)
		if err != nil {
			return nil, err
		}
		files = append(files, *f)
	}
	return files, rows.Err()
}

func (s *Store) FileCount() (int, error) {
	var count int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM local_files`).Scan(&count)
	return count, err
}

func (s *Store) NextRevision() (int64, error) {
	var rev int64
	err := s.db.QueryRow(`SELECT COALESCE(MAX(revision), 0) + 1 FROM local_files`).Scan(&rev)
	return rev, err
}

func (s *Store) MaxRevision() (int64, error) {
	var rev int64
	err := s.db.QueryRow(`SELECT COALESCE(MAX(revision), 0) FROM local_files`).Scan(&rev)
	return rev, err
}

func (s *Store) SaveSession(sess PlaybackSession) error {
	_, err := s.db.Exec(`
		INSERT OR REPLACE INTO playback_sessions
		(session_id, token_hash, local_file_id, idle_timeout_seconds, absolute_expires_at, last_activity)
		VALUES (?, ?, ?, ?, ?, ?)
	`, sess.SessionID, sess.TokenHash, sess.LocalFileID, sess.IdleTimeoutSeconds,
		sess.AbsoluteExpiresAt.UTC().Format(time.RFC3339), sess.LastActivity.UTC().Format(time.RFC3339))
	return err
}

func (s *Store) GetSessionByTokenHash(tokenHash string) (*PlaybackSession, error) {
	row := s.db.QueryRow(`
		SELECT session_id, token_hash, local_file_id, idle_timeout_seconds, absolute_expires_at, last_activity
		FROM playback_sessions WHERE token_hash = ?
	`, tokenHash)
	return scanSession(row)
}

func (s *Store) DeleteSession(sessionID string) error {
	_, err := s.db.Exec(`DELETE FROM playback_sessions WHERE session_id = ?`, sessionID)
	return err
}

func (s *Store) TouchSession(sessionID string) error {
	_, err := s.db.Exec(`UPDATE playback_sessions SET last_activity = ? WHERE session_id = ?`,
		time.Now().UTC().Format(time.RFC3339), sessionID)
	return err
}

func (s *Store) ActiveSessions() ([]PlaybackSession, error) {
	rows, err := s.db.Query(`
		SELECT session_id, token_hash, local_file_id, idle_timeout_seconds, absolute_expires_at, last_activity
		FROM playback_sessions
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var sessions []PlaybackSession
	for rows.Next() {
		sess, err := scanSession(rows)
		if err != nil {
			return nil, err
		}
		sessions = append(sessions, *sess)
	}
	return sessions, rows.Err()
}

func (s *Store) SetState(key, value string) error {
	_, err := s.db.Exec(`
		INSERT INTO node_state (key, value) VALUES (?, ?)
		ON CONFLICT(key) DO UPDATE SET value=excluded.value
	`, key, value)
	return err
}

func (s *Store) GetState(key string) (string, error) {
	var value string
	err := s.db.QueryRow(`SELECT value FROM node_state WHERE key = ?`, key).Scan(&value)
	if err == sql.ErrNoRows {
		return "", nil
	}
	return value, err
}

func FileID(path string) string {
	sum := sha256.Sum256([]byte(path))
	return fmt.Sprintf("f-%s", hex.EncodeToString(sum[:])[:12])
}

type scannable interface {
	Scan(dest ...any) error
}

func scanFile(row scannable) (*LocalFile, error) {
	var f LocalFile
	var modStr string
	if err := row.Scan(&f.LocalFileID, &f.Path, &f.SizeBytes, &modStr, &f.MediaType, &f.ReleaseName, &f.Revision); err != nil {
		return nil, err
	}
	f.ModifiedAt, _ = time.Parse(time.RFC3339, modStr)
	return &f, nil
}

func scanSession(row scannable) (*PlaybackSession, error) {
	var s PlaybackSession
	var absStr, lastStr string
	if err := row.Scan(&s.SessionID, &s.TokenHash, &s.LocalFileID, &s.IdleTimeoutSeconds, &absStr, &lastStr); err != nil {
		return nil, err
	}
	s.AbsoluteExpiresAt, _ = time.Parse(time.RFC3339, absStr)
	s.LastActivity, _ = time.Parse(time.RFC3339, lastStr)
	return &s, nil
}