Warp/server/internal/storage/sqlstore/media.go

package sqlstore

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Media is the shared SQL implementation of storage.MediaStore.
type Media struct {
	db       *sql.DB
	rebind   func(string) string
	isUnique func(error) bool
}

// NewMedia builds a media store.
func NewMedia(db *sql.DB, rebind func(string) string, isUnique func(error) bool) *Media {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	if isUnique == nil {
		isUnique = func(error) bool { return false }
	}
	return &Media{db: db, rebind: rebind, isUnique: isUnique}
}

func (s *Media) q(query string) string { return s.rebind(query) }

const insertMediaSQL = `
	INSERT INTO media_repository
		(media_origin, media_id, content_type, upload_name, file_size, sha256,
		 user_id, created_at_ms, pending, expires_at_ms)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

// CreateMedia implements storage.MediaStore.
func (s *Media) CreateMedia(ctx context.Context, m storage.Media) error {
	_, err := s.db.ExecContext(ctx, s.q(insertMediaSQL),
		m.Origin, m.MediaID, m.ContentType, m.UploadName, m.Size, m.SHA256,
		m.UserID, m.CreatedAtMS, m.Pending, m.ExpiresAtMS)
	switch {
	case err == nil:
		return nil
	case s.isUnique(err):
		return storage.ErrMediaExists
	default:
		return fmt.Errorf("store media %s/%s: %w", m.Origin, m.MediaID, err)
	}
}

const selectMediaSQL = `
	SELECT media_origin, media_id, content_type, upload_name, file_size, sha256,
		user_id, created_at_ms, pending, expires_at_ms
	FROM media_repository WHERE media_origin = ? AND media_id = ?`

// Media implements storage.MediaStore.
func (s *Media) Media(ctx context.Context, origin, mediaID string) (storage.Media, error) {
	var m storage.Media
	err := s.db.QueryRowContext(ctx, s.q(selectMediaSQL), origin, mediaID).Scan(
		&m.Origin, &m.MediaID, &m.ContentType, &m.UploadName, &m.Size, &m.SHA256,
		&m.UserID, &m.CreatedAtMS, &m.Pending, &m.ExpiresAtMS)
	switch {
	case err == nil:
		return m, nil
	case errors.Is(err, sql.ErrNoRows):
		return storage.Media{}, storage.ErrNoSuchMedia
	default:
		return storage.Media{}, fmt.Errorf("read media %s/%s: %w", origin, mediaID, err)
	}
}

// completeUploadSQL only touches pending media, so that two uploads racing
// for the same ID cannot both win: the second matches no row.
const completeUploadSQL = `
	UPDATE media_repository
	SET content_type = ?, upload_name = ?, file_size = ?, sha256 = ?,
		pending = ?, expires_at_ms = 0
	WHERE media_origin = ? AND media_id = ? AND pending = ?`

// CompleteUpload implements storage.MediaStore.
func (s *Media) CompleteUpload(ctx context.Context, origin, mediaID string, c storage.MediaContent) error {
	res, err := s.db.ExecContext(ctx, s.q(completeUploadSQL),
		c.ContentType, c.UploadName, c.Size, c.SHA256, false, origin, mediaID, true)
	if err != nil {
		return fmt.Errorf("complete upload %s/%s: %w", origin, mediaID, err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("complete upload %s/%s: %w", origin, mediaID, err)
	}
	if n == 1 {
		return nil
	}
	// Nothing pending matched: either there is no such media, or its
	// content is already there.
	if _, err := s.Media(ctx, origin, mediaID); err != nil {
		return err
	}
	return storage.ErrMediaUploaded
}

const countPendingSQL = `
	SELECT COUNT(*) FROM media_repository
	WHERE pending = ? AND user_id = ? AND expires_at_ms > ?`

// PendingUploads implements storage.MediaStore.
func (s *Media) PendingUploads(ctx context.Context, userID string, nowMS int64) (int, error) {
	var n int
	if err := s.db.QueryRowContext(ctx, s.q(countPendingSQL), true, userID, nowMS).Scan(&n); err != nil {
		return 0, fmt.Errorf("count the pending uploads of %s: %w", userID, err)
	}
	return n, nil
}

const deleteExpiredSQL = `
	DELETE FROM media_repository WHERE pending = ? AND expires_at_ms <= ?`

// DeleteExpiredUploads implements storage.MediaStore.
func (s *Media) DeleteExpiredUploads(ctx context.Context, nowMS int64) (int, error) {
	res, err := s.db.ExecContext(ctx, s.q(deleteExpiredSQL), true, nowMS)
	if err != nil {
		return 0, fmt.Errorf("delete expired uploads: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("delete expired uploads: %w", err)
	}
	return int(n), nil
}

const selectThumbnailSQL = `
	SELECT media_origin, media_id, width, height, method, content_type,
		file_size, sha256, created_at_ms
	FROM media_thumbnails
	WHERE media_origin = ? AND media_id = ? AND width = ? AND height = ? AND method = ?`

// Thumbnail implements storage.MediaStore.
func (s *Media) Thumbnail(ctx context.Context, origin, mediaID string, width, height int, method string) (storage.Thumbnail, error) {
	var t storage.Thumbnail
	err := s.db.QueryRowContext(ctx, s.q(selectThumbnailSQL), origin, mediaID, width, height, method).Scan(
		&t.Origin, &t.MediaID, &t.Width, &t.Height, &t.Method, &t.ContentType,
		&t.Size, &t.SHA256, &t.CreatedAtMS)
	switch {
	case err == nil:
		return t, nil
	case errors.Is(err, sql.ErrNoRows):
		return storage.Thumbnail{}, storage.ErrNoSuchMedia
	default:
		return storage.Thumbnail{}, fmt.Errorf("read a thumbnail of %s/%s: %w", origin, mediaID, err)
	}
}

const upsertThumbnailSQL = `
	INSERT INTO media_thumbnails
		(media_origin, media_id, width, height, method, content_type,
		 file_size, sha256, created_at_ms)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT (media_origin, media_id, width, height, method) DO UPDATE SET
		content_type = excluded.content_type,
		file_size = excluded.file_size,
		sha256 = excluded.sha256,
		created_at_ms = excluded.created_at_ms`

// StoreThumbnail implements storage.MediaStore.
func (s *Media) StoreThumbnail(ctx context.Context, t storage.Thumbnail) error {
	if _, err := s.db.ExecContext(ctx, s.q(upsertThumbnailSQL),
		t.Origin, t.MediaID, t.Width, t.Height, t.Method, t.ContentType,
		t.Size, t.SHA256, t.CreatedAtMS); err != nil {
		return fmt.Errorf("store a thumbnail of %s/%s: %w", t.Origin, t.MediaID, err)
	}
	return nil
}

const selectMediaUsageSQL = `
	SELECT COUNT(*), COALESCE(SUM(file_size), 0) FROM media_repository
	WHERE user_id = ? AND media_origin = ? AND pending = ?`

// MediaUsage implements storage.MediaStore.
//
// It counts every upload at its own size, even when two of them are the
// same bytes and the blob store holds one copy: a quota is about what a
// user asked this server to keep, and how the server keeps it is its own
// business.
func (s *Media) MediaUsage(ctx context.Context, origin, userID string) (storage.MediaUsage, error) {
	var u storage.MediaUsage
	if err := s.db.QueryRowContext(ctx, s.q(selectMediaUsageSQL), userID, origin, false).
		Scan(&u.Count, &u.Bytes); err != nil {
		return storage.MediaUsage{}, fmt.Errorf("total the uploads of %s: %w", userID, err)
	}
	return u, nil
}

const (
	selectMediaQuotaSQL = `SELECT max_bytes FROM media_quotas WHERE user_id = ?`

	upsertMediaQuotaSQL = `
		INSERT INTO media_quotas (user_id, max_bytes) VALUES (?, ?)
		ON CONFLICT (user_id) DO UPDATE SET max_bytes = excluded.max_bytes`

	deleteMediaQuotaSQL = `DELETE FROM media_quotas WHERE user_id = ?`
)

// MediaQuota implements storage.MediaStore.
func (s *Media) MediaQuota(ctx context.Context, userID string) (int64, bool, error) {
	var maxBytes int64
	err := s.db.QueryRowContext(ctx, s.q(selectMediaQuotaSQL), userID).Scan(&maxBytes)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return 0, false, nil
	case err != nil:
		return 0, false, fmt.Errorf("read the media quota of %s: %w", userID, err)
	}
	return maxBytes, true, nil
}

// SetMediaQuota implements storage.MediaStore.
func (s *Media) SetMediaQuota(ctx context.Context, userID string, maxBytes int64) error {
	if _, err := s.db.ExecContext(ctx, s.q(upsertMediaQuotaSQL), userID, maxBytes); err != nil {
		return fmt.Errorf("set the media quota of %s: %w", userID, err)
	}
	return nil
}

// DeleteMediaQuota implements storage.MediaStore.
func (s *Media) DeleteMediaQuota(ctx context.Context, userID string) error {
	if _, err := s.db.ExecContext(ctx, s.q(deleteMediaQuotaSQL), userID); err != nil {
		return fmt.Errorf("delete the media quota of %s: %w", userID, err)
	}
	return nil
}