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

package pocketbase

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

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

// Media is the PocketBase implementation of storage.MediaStore: the
// metadata of media. The content is the blob store's.
type Media struct{ *base }

// CreateMedia implements storage.MediaStore.
func (s *Media) CreateMedia(ctx context.Context, m storage.Media) error {
	return s.write(ctx, func(t *tx) error {
		taken, err := t.exists("media_repository", `media_origin = ? AND media_id = ?`, m.Origin, m.MediaID)
		if err != nil {
			return fmt.Errorf("store media %s/%s: %w", m.Origin, m.MediaID, err)
		}
		if taken {
			return storage.ErrMediaExists
		}
		t.put("media_repository", mediaRow(m))
		return nil
	})
}

func mediaRow(m storage.Media) rec {
	return rec{
		"media_origin": m.Origin, "media_id": m.MediaID, "content_type": m.ContentType,
		"upload_name": m.UploadName, "file_size": m.Size, "sha256": m.SHA256,
		"user_id": m.UserID, "created_at_ms": m.CreatedAtMS, "pending": m.Pending,
		"expires_at_ms": m.ExpiresAtMS,
	}
}

// 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, `
		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 = ?`, 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)
	}
}

// CompleteUpload implements storage.MediaStore. Only pending media is
// completed, so that two uploads for the same ID cannot both win.
func (s *Media) CompleteUpload(ctx context.Context, origin, mediaID string, c storage.MediaContent) error {
	var completed bool
	err := s.write(ctx, func(t *tx) error {
		n, err := t.setWhere("media_repository", rec{
			"content_type": c.ContentType, "upload_name": c.UploadName, "file_size": c.Size,
			"sha256": c.SHA256, "pending": false, "expires_at_ms": int64(0),
		}, `media_origin = ? AND media_id = ? AND pending = ?`, origin, mediaID, true)
		completed = n == 1
		return err
	})
	if err != nil {
		return fmt.Errorf("complete upload %s/%s: %w", origin, mediaID, err)
	}
	if completed {
		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
}

// 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, `SELECT COUNT(*) FROM media_repository
		WHERE pending = ? AND user_id = ? AND expires_at_ms > ?`, true, userID, nowMS).Scan(&n); err != nil {
		return 0, fmt.Errorf("count the pending uploads of %s: %w", userID, err)
	}
	return n, nil
}

// DeleteExpiredUploads implements storage.MediaStore.
func (s *Media) DeleteExpiredUploads(ctx context.Context, nowMS int64) (int, error) {
	var n int
	err := s.write(ctx, func(t *tx) error {
		var err error
		n, err = t.delWhere("media_repository", `pending = ? AND expires_at_ms <= ?`, true, nowMS)
		return err
	})
	if err != nil {
		return 0, fmt.Errorf("delete expired uploads: %w", err)
	}
	return n, nil
}

// 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, `
		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 = ?`,
		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)
	}
}

// StoreThumbnail implements storage.MediaStore.
func (s *Media) StoreThumbnail(ctx context.Context, t storage.Thumbnail) error {
	err := s.write(ctx, func(w *tx) error {
		w.put("media_thumbnails", rec{
			"media_origin": t.Origin, "media_id": t.MediaID, "width": t.Width, "height": t.Height,
			"method": t.Method, "content_type": t.ContentType, "file_size": t.Size,
			"sha256": t.SHA256, "created_at_ms": t.CreatedAtMS,
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("store a thumbnail of %s/%s: %w", t.Origin, t.MediaID, err)
	}
	return nil
}

// MediaUsage implements storage.MediaStore. It counts every upload at its
// own size, even when two of them are the same bytes.
func (s *Media) MediaUsage(ctx context.Context, origin, userID string) (storage.MediaUsage, error) {
	var u storage.MediaUsage
	if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(file_size), 0) FROM media_repository
		WHERE user_id = ? AND media_origin = ? AND pending = ?`, 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
}

// MediaQuota implements storage.MediaStore.
func (s *Media) MediaQuota(ctx context.Context, userID string) (int64, bool, error) {
	var maxBytes int64
	err := s.db.QueryRowContext(ctx, `SELECT max_bytes FROM media_quotas WHERE user_id = ?`, 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 {
	err := s.write(ctx, func(t *tx) error {
		t.put("media_quotas", rec{"user_id": userID, "max_bytes": maxBytes})
		return nil
	})
	if 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 {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delKey("media_quotas", userID)
		return err
	})
	if err != nil {
		return fmt.Errorf("delete the media quota of %s: %w", userID, err)
	}
	return nil
}

var _ storage.MediaStore = (*Media)(nil)