stash/pkg/sqlite/scene_marker.go

package sqlite

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

	"github.com/doug-martin/goqu/v9"
	"github.com/doug-martin/goqu/v9/exp"
	"github.com/jmoiron/sqlx"
	"gopkg.in/guregu/null.v4"

	"github.com/stashapp/stash/pkg/models"
)

const (
	sceneMarkerTable      = "scene_markers"
	sceneMarkersTagsTable = "scene_markers_tags"
	sceneMarkerIDColumn   = "scene_marker_id"
)

const countSceneMarkersForTagQuery = `
SELECT scene_markers.id FROM scene_markers
LEFT JOIN scene_markers_tags as tags_join on tags_join.scene_marker_id = scene_markers.id
WHERE tags_join.tag_id = ? OR scene_markers.primary_tag_id = ?
GROUP BY scene_markers.id
`

type sceneMarkerRow struct {
	ID           int        `db:"id" goqu:"skipinsert"`
	Title        string     `db:"title"` // TODO: make db schema (and gql schema) nullable
	Seconds      float64    `db:"seconds"`
	PrimaryTagID int        `db:"primary_tag_id"`
	SceneID      int        `db:"scene_id"`
	CreatedAt    Timestamp  `db:"created_at"`
	UpdatedAt    Timestamp  `db:"updated_at"`
	EndSeconds   null.Float `db:"end_seconds"`
}

func (r *sceneMarkerRow) fromSceneMarker(o models.SceneMarker) {
	r.ID = o.ID
	r.Title = o.Title
	r.Seconds = o.Seconds
	if o.EndSeconds != nil {
		r.EndSeconds = null.FloatFrom(*o.EndSeconds)
	}
	r.PrimaryTagID = o.PrimaryTagID
	r.SceneID = o.SceneID
	r.CreatedAt = Timestamp{Timestamp: o.CreatedAt}
	r.UpdatedAt = Timestamp{Timestamp: o.UpdatedAt}
}

func (r *sceneMarkerRow) resolve() *models.SceneMarker {
	ret := &models.SceneMarker{
		ID:           r.ID,
		Title:        r.Title,
		Seconds:      r.Seconds,
		EndSeconds:   r.EndSeconds.Ptr(),
		PrimaryTagID: r.PrimaryTagID,
		SceneID:      r.SceneID,
		CreatedAt:    r.CreatedAt.Timestamp,
		UpdatedAt:    r.UpdatedAt.Timestamp,
	}

	return ret
}

type sceneMarkerRowRecord struct {
	updateRecord
}

func (r *sceneMarkerRowRecord) fromPartial(o models.SceneMarkerPartial) {
	// TODO: replace with setNullString after schema is made nullable
	// r.setNullString("title", o.Title)
	// saves a null input as the empty string
	if o.Title.Set {
		r.set("title", o.Title.Value)
	}
	r.setFloat64("seconds", o.Seconds)
	r.setNullFloat64("end_seconds", o.EndSeconds)
	r.setInt("primary_tag_id", o.PrimaryTagID)
	r.setInt("scene_id", o.SceneID)
	r.setTimestamp("created_at", o.CreatedAt)
	r.setTimestamp("updated_at", o.UpdatedAt)
}

type sceneMarkerRepositoryType struct {
	repository

	scenes repository
	tags   joinRepository
}

var (
	sceneMarkerRepository = sceneMarkerRepositoryType{
		repository: repository{
			tableName: sceneMarkerTable,
			idColumn:  idColumn,
		},
		scenes: repository{
			tableName: sceneTable,
			idColumn:  idColumn,
		},
		tags: joinRepository{
			repository: repository{
				tableName: sceneMarkersTagsTable,
				idColumn:  sceneMarkerIDColumn,
			},
			fkColumn: tagIDColumn,
		},
	}
)

type SceneMarkerStore struct{}

func NewSceneMarkerStore() *SceneMarkerStore {
	return &SceneMarkerStore{}
}

func (qb *SceneMarkerStore) table() exp.IdentifierExpression {
	return sceneMarkerTableMgr.table
}

func (qb *SceneMarkerStore) selectDataset() *goqu.SelectDataset {
	return dialect.From(qb.table()).Select(qb.table().All())
}

func (qb *SceneMarkerStore) Create(ctx context.Context, newObject *models.SceneMarker) error {
	var r sceneMarkerRow
	r.fromSceneMarker(*newObject)

	id, err := sceneMarkerTableMgr.insertID(ctx, r)
	if err != nil {
		return err
	}

	updated, err := qb.find(ctx, id)
	if err != nil {
		return fmt.Errorf("finding after create: %w", err)
	}

	*newObject = *updated

	return nil
}

func (qb *SceneMarkerStore) UpdatePartial(ctx context.Context, id int, partial models.SceneMarkerPartial) (*models.SceneMarker, error) {
	r := sceneMarkerRowRecord{
		updateRecord{
			Record: make(exp.Record),
		},
	}

	r.fromPartial(partial)

	if len(r.Record) > 0 {
		if err := sceneMarkerTableMgr.updateByID(ctx, id, r.Record); err != nil {
			return nil, err
		}
	}

	if partial.TagIDs != nil {
		if err := sceneMarkersTagsTableMgr.modifyJoins(ctx, id, partial.TagIDs.IDs, partial.TagIDs.Mode); err != nil {
			return nil, fmt.Errorf("modifying scene marker tags: %w", err)
		}
	}

	return qb.find(ctx, id)
}

func (qb *SceneMarkerStore) Update(ctx context.Context, updatedObject *models.SceneMarker) error {
	var r sceneMarkerRow
	r.fromSceneMarker(*updatedObject)

	if err := sceneMarkerTableMgr.updateByID(ctx, updatedObject.ID, r); err != nil {
		return err
	}

	return nil
}

func (qb *SceneMarkerStore) Destroy(ctx context.Context, id int) error {
	return sceneMarkerRepository.destroyExisting(ctx, []int{id})
}

// returns nil, nil if not found
func (qb *SceneMarkerStore) Find(ctx context.Context, id int) (*models.SceneMarker, error) {
	ret, err := qb.find(ctx, id)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, nil
	}
	return ret, err
}

func (qb *SceneMarkerStore) FindMany(ctx context.Context, ids []int) ([]*models.SceneMarker, error) {
	ret := make([]*models.SceneMarker, len(ids))

	table := qb.table()
	q := qb.selectDataset().Prepared(true).Where(table.Col(idColumn).In(ids))
	unsorted, err := qb.getMany(ctx, q)
	if err != nil {
		return nil, err
	}

	for _, s := range unsorted {
		i := slices.Index(ids, s.ID)
		ret[i] = s
	}

	for i := range ret {
		if ret[i] == nil {
			return nil, fmt.Errorf("scene marker with id %d not found", ids[i])
		}
	}

	return ret, nil
}

// returns nil, sql.ErrNoRows if not found
func (qb *SceneMarkerStore) find(ctx context.Context, id int) (*models.SceneMarker, error) {
	q := qb.selectDataset().Where(sceneMarkerTableMgr.byID(id))

	ret, err := qb.get(ctx, q)
	if err != nil {
		return nil, err
	}

	return ret, nil
}

// returns nil, sql.ErrNoRows if not found
func (qb *SceneMarkerStore) get(ctx context.Context, q *goqu.SelectDataset) (*models.SceneMarker, error) {
	ret, err := qb.getMany(ctx, q)
	if err != nil {
		return nil, err
	}

	if len(ret) == 0 {
		return nil, sql.ErrNoRows
	}

	return ret[0], nil
}

func (qb *SceneMarkerStore) getMany(ctx context.Context, q *goqu.SelectDataset) ([]*models.SceneMarker, error) {
	const single = false
	var ret []*models.SceneMarker
	if err := queryFunc(ctx, q, single, func(r *sqlx.Rows) error {
		var f sceneMarkerRow
		if err := r.StructScan(&f); err != nil {
			return err
		}

		s := f.resolve()

		ret = append(ret, s)
		return nil
	}); err != nil {
		return nil, err
	}

	return ret, nil
}

func (qb *SceneMarkerStore) FindBySceneID(ctx context.Context, sceneID int) ([]*models.SceneMarker, error) {
	query := `
		SELECT scene_markers.* FROM scene_markers
		WHERE scene_markers.scene_id = ?
		GROUP BY scene_markers.id
		ORDER BY scene_markers.seconds ASC
	`
	args := []interface{}{sceneID}
	return qb.querySceneMarkers(ctx, query, args)
}

func (qb *SceneMarkerStore) CountByTagID(ctx context.Context, tagID int) (int, error) {
	args := []interface{}{tagID, tagID}
	return sceneMarkerRepository.runCountQuery(ctx, sceneMarkerRepository.buildCountQuery(countSceneMarkersForTagQuery), args)
}

func (qb *SceneMarkerStore) GetMarkerStrings(ctx context.Context, q *string, sort *string) ([]*models.MarkerStringsResultType, error) {
	query := "SELECT count(*) as `count`, scene_markers.id as id, scene_markers.title as title FROM scene_markers"
	if q != nil {
		query += " WHERE title LIKE '%" + *q + "%'"
	}
	query += " GROUP BY title"
	if sort != nil && *sort == "count" {
		query += " ORDER BY `count` DESC"
	} else {
		query += " ORDER BY title ASC"
	}
	var args []interface{}
	return qb.queryMarkerStringsResultType(ctx, query, args)
}

func (qb *SceneMarkerStore) Wall(ctx context.Context, q *string) ([]*models.SceneMarker, error) {
	s := ""
	if q != nil {
		s = *q
	}

	table := qb.table()
	qq := qb.selectDataset().Prepared(true).Where(table.Col("title").Like("%" + s + "%")).Order(goqu.L("RANDOM()").Asc()).Limit(80)
	return qb.getMany(ctx, qq)
}

func (qb *SceneMarkerStore) makeQuery(ctx context.Context, sceneMarkerFilter *models.SceneMarkerFilterType, findFilter *models.FindFilterType) (*queryBuilder, error) {
	if sceneMarkerFilter == nil {
		sceneMarkerFilter = &models.SceneMarkerFilterType{}
	}
	if findFilter == nil {
		findFilter = &models.FindFilterType{}
	}

	query := sceneMarkerRepository.newQuery()
	distinctIDs(&query, sceneMarkerTable)

	if q := findFilter.Q; q != nil && *q != "" {
		query.join(sceneTable, "", "scenes.id = scene_markers.scene_id")
		query.join(tagTable, "", "scene_markers.primary_tag_id = tags.id")
		searchColumns := []string{"scene_markers.title", "scenes.title", "tags.name"}
		query.parseQueryString(searchColumns, *q)
	}

	filter := filterBuilderFromHandler(ctx, &sceneMarkerFilterHandler{
		sceneMarkerFilter: sceneMarkerFilter,
	})

	if err := query.addFilter(filter); err != nil {
		return nil, err
	}

	if err := qb.setSceneMarkerSort(&query, findFilter); err != nil {
		return nil, err
	}
	query.sortAndPagination += getPagination(findFilter)

	return &query, nil
}

func (qb *SceneMarkerStore) Query(ctx context.Context, sceneMarkerFilter *models.SceneMarkerFilterType, findFilter *models.FindFilterType) ([]*models.SceneMarker, int, error) {
	query, err := qb.makeQuery(ctx, sceneMarkerFilter, findFilter)
	if err != nil {
		return nil, 0, err
	}

	idsResult, countResult, err := query.executeFind(ctx)
	if err != nil {
		return nil, 0, err
	}

	sceneMarkers, err := qb.FindMany(ctx, idsResult)
	if err != nil {
		return nil, 0, err
	}

	return sceneMarkers, countResult, nil
}

func (qb *SceneMarkerStore) QueryCount(ctx context.Context, sceneMarkerFilter *models.SceneMarkerFilterType, findFilter *models.FindFilterType) (int, error) {
	query, err := qb.makeQuery(ctx, sceneMarkerFilter, findFilter)
	if err != nil {
		return 0, err
	}

	return query.executeCount(ctx)
}

var sceneMarkerSortOptions = sortOptions{
	"created_at",
	"id",
	"title",
	"random",
	"scene_id",
	"scenes_updated_at",
	"seconds",
	"updated_at",
	"duration",
}

func (qb *SceneMarkerStore) setSceneMarkerSort(query *queryBuilder, findFilter *models.FindFilterType) error {
	sort := findFilter.GetSort("title")
	direction := findFilter.GetDirection()

	// CVE-2024-32231 - ensure sort is in the list of allowed sorts
	if err := sceneMarkerSortOptions.validateSort(sort); err != nil {
		return err
	}

	switch sort {
	case "scenes_updated_at":
		sort = "updated_at"
		query.joinSort(sceneTable, "", "scenes.id = scene_markers.scene_id")
		query.sortAndPagination += getSort(sort, direction, sceneTable)
	case "title":
		query.joinSort(tagTable, "", "scene_markers.primary_tag_id = tags.id")
		query.sortAndPagination += " ORDER BY COALESCE(NULLIF(scene_markers.title,''), tags.name) COLLATE NATURAL_CI " + direction
	case "duration":
		sort = "(scene_markers.end_seconds - scene_markers.seconds)"
		query.sortAndPagination += getSort(sort, direction, sceneMarkerTable)
	default:
		query.sortAndPagination += getSort(sort, direction, sceneMarkerTable)
	}

	query.sortAndPagination += ", scene_markers.scene_id ASC, scene_markers.seconds ASC"
	return nil
}

func (qb *SceneMarkerStore) querySceneMarkers(ctx context.Context, query string, args []interface{}) ([]*models.SceneMarker, error) {
	const single = false
	var ret []*models.SceneMarker
	if err := sceneMarkerRepository.queryFunc(ctx, query, args, single, func(r *sqlx.Rows) error {
		var f sceneMarkerRow
		if err := r.StructScan(&f); err != nil {
			return err
		}

		s := f.resolve()

		ret = append(ret, s)
		return nil
	}); err != nil {
		return nil, err
	}

	return ret, nil
}

func (qb *SceneMarkerStore) queryMarkerStringsResultType(ctx context.Context, query string, args []interface{}) ([]*models.MarkerStringsResultType, error) {
	rows, err := dbWrapper.Queryx(ctx, query, args...)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, err
	}
	defer rows.Close()

	markerStrings := make([]*models.MarkerStringsResultType, 0)
	for rows.Next() {
		markerString := models.MarkerStringsResultType{}
		if err := rows.StructScan(&markerString); err != nil {
			return nil, err
		}
		markerStrings = append(markerStrings, &markerString)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return markerStrings, nil
}

func (qb *SceneMarkerStore) GetTagIDs(ctx context.Context, id int) ([]int, error) {
	return sceneMarkerRepository.tags.getIDs(ctx, id)
}

func (qb *SceneMarkerStore) UpdateTags(ctx context.Context, id int, tagIDs []int) error {
	// Delete the existing joins and then create new ones
	return sceneMarkerRepository.tags.replace(ctx, id, tagIDs)
}

func (qb *SceneMarkerStore) Count(ctx context.Context) (int, error) {
	q := dialect.Select(goqu.COUNT("*")).From(qb.table())
	return count(ctx, q)
}

func (qb *SceneMarkerStore) All(ctx context.Context) ([]*models.SceneMarker, error) {
	return qb.getMany(ctx, qb.selectDataset())
}