stash/pkg/sqlite/folder.go

package sqlite

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

	"github.com/doug-martin/goqu/v9"
	"github.com/doug-martin/goqu/v9/exp"
	"github.com/jmoiron/sqlx"
	"github.com/stashapp/stash/pkg/models"
	"gopkg.in/guregu/null.v4"
)

const folderTable = "folders"
const folderIDColumn = "folder_id"

type folderRow struct {
	ID             models.FolderID `db:"id" goqu:"skipinsert"`
	Path           string          `db:"path"`
	ZipFileID      null.Int        `db:"zip_file_id"`
	ParentFolderID null.Int        `db:"parent_folder_id"`
	ModTime        Timestamp       `db:"mod_time"`
	CreatedAt      Timestamp       `db:"created_at"`
	UpdatedAt      Timestamp       `db:"updated_at"`
}

func (r *folderRow) fromFolder(o models.Folder) {
	r.ID = o.ID
	r.Path = o.Path
	r.ZipFileID = nullIntFromFileIDPtr(o.ZipFileID)
	r.ParentFolderID = nullIntFromFolderIDPtr(o.ParentFolderID)
	r.ModTime = Timestamp{Timestamp: o.ModTime}
	r.CreatedAt = Timestamp{Timestamp: o.CreatedAt}
	r.UpdatedAt = Timestamp{Timestamp: o.UpdatedAt}
}

type folderQueryRow struct {
	folderRow

	ZipBasename   null.String `db:"zip_basename"`
	ZipFolderPath null.String `db:"zip_folder_path"`
	ZipSize       null.Int    `db:"zip_size"`
}

func (r *folderQueryRow) resolve() *models.Folder {
	ret := &models.Folder{
		ID: r.ID,
		DirEntry: models.DirEntry{
			ZipFileID: nullIntFileIDPtr(r.ZipFileID),
			ModTime:   r.ModTime.Timestamp,
		},
		Path:           string(r.Path),
		ParentFolderID: nullIntFolderIDPtr(r.ParentFolderID),
		CreatedAt:      r.CreatedAt.Timestamp,
		UpdatedAt:      r.UpdatedAt.Timestamp,
	}

	if ret.ZipFileID != nil && r.ZipFolderPath.Valid && r.ZipBasename.Valid {
		ret.ZipFile = &models.BaseFile{
			ID:       *ret.ZipFileID,
			Path:     filepath.Join(r.ZipFolderPath.String, r.ZipBasename.String),
			Basename: r.ZipBasename.String,
			Size:     r.ZipSize.Int64,
		}
	}

	return ret
}

type folderQueryRows []folderQueryRow

func (r folderQueryRows) resolve() []*models.Folder {
	var ret []*models.Folder

	for _, row := range r {
		f := row.resolve()
		ret = append(ret, f)
	}

	return ret
}

type folderRepositoryType struct {
	repository

	galleries repository
}

var (
	folderRepository = folderRepositoryType{
		repository: repository{
			tableName: folderTable,
			idColumn:  idColumn,
		},
		galleries: repository{
			tableName: galleryTable,
			idColumn:  folderIDColumn,
		},
	}
)

type FolderStore struct {
	repository

	tableMgr *table
}

func NewFolderStore() *FolderStore {
	return &FolderStore{
		repository: repository{
			tableName: folderTable,
			idColumn:  idColumn,
		},

		tableMgr: folderTableMgr,
	}
}

func (qb *FolderStore) Create(ctx context.Context, f *models.Folder) error {
	var r folderRow
	r.fromFolder(*f)

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

	// only assign id once we are successful
	f.ID = models.FolderID(id)

	return nil
}

func (qb *FolderStore) Update(ctx context.Context, updatedObject *models.Folder) error {
	var r folderRow
	r.fromFolder(*updatedObject)

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

	return nil
}

func (qb *FolderStore) Destroy(ctx context.Context, id models.FolderID) error {
	return qb.tableMgr.destroyExisting(ctx, []int{int(id)})
}

func (qb *FolderStore) table() exp.IdentifierExpression {
	return qb.tableMgr.table
}

func (qb *FolderStore) selectDataset() *goqu.SelectDataset {
	table := qb.table()
	fileTable := fileTableMgr.table

	zipFileTable := fileTable.As("zip_files")
	zipFolderTable := table.As("zip_files_folders")

	cols := []interface{}{
		table.Col("id"),
		table.Col("path"),
		table.Col("zip_file_id"),
		table.Col("parent_folder_id"),
		table.Col("mod_time"),
		table.Col("created_at"),
		table.Col("updated_at"),
		zipFileTable.Col("basename").As("zip_basename"),
		zipFolderTable.Col("path").As("zip_folder_path"),
		// size is needed to open containing zip files
		zipFileTable.Col("size").As("zip_size"),
	}

	ret := dialect.From(table).Select(cols...)

	return ret.LeftJoin(
		zipFileTable,
		goqu.On(table.Col("zip_file_id").Eq(zipFileTable.Col("id"))),
	).LeftJoin(
		zipFolderTable,
		goqu.On(zipFileTable.Col("parent_folder_id").Eq(zipFolderTable.Col(idColumn))),
	)
}

func (qb *FolderStore) countDataset() *goqu.SelectDataset {
	table := qb.table()
	fileTable := fileTableMgr.table

	zipFileTable := fileTable.As("zip_files")
	zipFolderTable := table.As("zip_files_folders")

	ret := dialect.From(table).Select(goqu.COUNT(goqu.DISTINCT(table.Col("id"))))

	return ret.LeftJoin(
		zipFileTable,
		goqu.On(table.Col("zip_file_id").Eq(zipFileTable.Col("id"))),
	).LeftJoin(
		zipFolderTable,
		goqu.On(zipFileTable.Col("parent_folder_id").Eq(zipFolderTable.Col(idColumn))),
	)
}

func (qb *FolderStore) get(ctx context.Context, q *goqu.SelectDataset) (*models.Folder, 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 *FolderStore) getMany(ctx context.Context, q *goqu.SelectDataset) ([]*models.Folder, error) {
	const single = false
	var rows folderQueryRows
	if err := queryFunc(ctx, q, single, func(r *sqlx.Rows) error {
		var f folderQueryRow
		if err := r.StructScan(&f); err != nil {
			return err
		}

		rows = append(rows, f)
		return nil
	}); err != nil {
		return nil, err
	}

	return rows.resolve(), nil
}

func (qb *FolderStore) Find(ctx context.Context, id models.FolderID) (*models.Folder, error) {
	q := qb.selectDataset().Where(qb.tableMgr.byID(id))

	ret, err := qb.get(ctx, q)
	if err != nil {
		return nil, fmt.Errorf("getting folder by id %d: %w", id, err)
	}

	return ret, nil
}

// FindByIDs finds multiple folders by their IDs.
// No check is made to see if the folders exist, and the order of the returned folders
// is not guaranteed to be the same as the order of the input IDs.
func (qb *FolderStore) FindByIDs(ctx context.Context, ids []models.FolderID) ([]*models.Folder, error) {
	folders := make([]*models.Folder, 0, len(ids))

	table := qb.table()
	if err := batchExec(ids, defaultBatchSize, func(batch []models.FolderID) error {
		q := qb.selectDataset().Prepared(true).Where(table.Col(idColumn).In(batch))
		unsorted, err := qb.getMany(ctx, q)
		if err != nil {
			return err
		}

		folders = append(folders, unsorted...)

		return nil
	}); err != nil {
		return nil, err
	}

	return folders, nil
}

func (qb *FolderStore) FindMany(ctx context.Context, ids []models.FolderID) ([]*models.Folder, error) {
	folders := make([]*models.Folder, len(ids))

	unsorted, err := qb.FindByIDs(ctx, ids)
	if err != nil {
		return nil, err
	}

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

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

	return folders, nil
}

func (qb *FolderStore) FindByPath(ctx context.Context, p string) (*models.Folder, error) {
	q := qb.selectDataset().Prepared(true).Where(qb.table().Col("path").Eq(p))

	ret, err := qb.get(ctx, q)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("getting folder by path %s: %w", p, err)
	}

	return ret, nil
}

func (qb *FolderStore) FindByParentFolderID(ctx context.Context, parentFolderID models.FolderID) ([]*models.Folder, error) {
	q := qb.selectDataset().Where(qb.table().Col("parent_folder_id").Eq(int(parentFolderID)))

	ret, err := qb.getMany(ctx, q)
	if err != nil {
		return nil, fmt.Errorf("getting folders by parent folder id %d: %w", parentFolderID, err)
	}

	return ret, nil
}

func (qb *FolderStore) allInPaths(q *goqu.SelectDataset, p []string) *goqu.SelectDataset {
	table := qb.table()

	var conds []exp.Expression
	for _, pp := range p {
		ppWildcard := pp + string(filepath.Separator) + "%"

		conds = append(conds, table.Col("path").Eq(pp), table.Col("path").Like(ppWildcard))
	}

	return q.Where(
		goqu.Or(conds...),
	)
}

// FindAllInPaths returns the all folders that are or are within any of the given paths.
// Returns all if limit is < 0.
// Returns all folders if p is empty.
func (qb *FolderStore) FindAllInPaths(ctx context.Context, p []string, limit, offset int) ([]*models.Folder, error) {
	q := qb.selectDataset().Prepared(true)
	q = qb.allInPaths(q, p)

	if limit > -1 {
		q = q.Limit(uint(limit))
	}

	q = q.Offset(uint(offset))

	ret, err := qb.getMany(ctx, q)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("getting folders in path %s: %w", p, err)
	}

	return ret, nil
}

// CountAllInPaths returns a count of all folders that are within any of the given paths.
// Returns count of all folders if p is empty.
func (qb *FolderStore) CountAllInPaths(ctx context.Context, p []string) (int, error) {
	q := qb.countDataset().Prepared(true)
	q = qb.allInPaths(q, p)

	return count(ctx, q)
}

// func (qb *FolderStore) findBySubquery(ctx context.Context, sq *goqu.SelectDataset) ([]*file.Folder, error) {
// 	table := qb.table()

// 	q := qb.selectDataset().Prepared(true).Where(
// 		table.Col(idColumn).Eq(
// 			sq,
// 		),
// 	)

// 	return qb.getMany(ctx, q)
// }

func (qb *FolderStore) FindByZipFileID(ctx context.Context, zipFileID models.FileID) ([]*models.Folder, error) {
	table := qb.table()

	q := qb.selectDataset().Prepared(true).Where(
		table.Col("zip_file_id").Eq(zipFileID),
	)

	return qb.getMany(ctx, q)
}

func (qb *FolderStore) validateFilter(fileFilter *models.FolderFilterType) error {
	const and = "AND"
	const or = "OR"
	const not = "NOT"

	if fileFilter.And != nil {
		if fileFilter.Or != nil {
			return illegalFilterCombination(and, or)
		}
		if fileFilter.Not != nil {
			return illegalFilterCombination(and, not)
		}

		return qb.validateFilter(fileFilter.And)
	}

	if fileFilter.Or != nil {
		if fileFilter.Not != nil {
			return illegalFilterCombination(or, not)
		}

		return qb.validateFilter(fileFilter.Or)
	}

	if fileFilter.Not != nil {
		return qb.validateFilter(fileFilter.Not)
	}

	return nil
}

func (qb *FolderStore) makeFilter(ctx context.Context, folderFilter *models.FolderFilterType) *filterBuilder {
	query := &filterBuilder{}

	if folderFilter.And != nil {
		query.and(qb.makeFilter(ctx, folderFilter.And))
	}
	if folderFilter.Or != nil {
		query.or(qb.makeFilter(ctx, folderFilter.Or))
	}
	if folderFilter.Not != nil {
		query.not(qb.makeFilter(ctx, folderFilter.Not))
	}

	filter := filterBuilderFromHandler(ctx, &folderFilterHandler{
		folderFilter: folderFilter,
	})

	return filter
}

func (qb *FolderStore) Query(ctx context.Context, options models.FolderQueryOptions) (*models.FolderQueryResult, error) {
	folderFilter := options.FolderFilter
	findFilter := options.FindFilter

	if folderFilter == nil {
		folderFilter = &models.FolderFilterType{}
	}
	if findFilter == nil {
		findFilter = &models.FindFilterType{}
	}

	query := qb.newQuery()

	distinctIDs(&query, folderTable)

	if q := findFilter.Q; q != nil && *q != "" {
		searchColumns := []string{"folders.path"}
		query.parseQueryString(searchColumns, *q)
	}

	if err := qb.validateFilter(folderFilter); err != nil {
		return nil, err
	}
	filter := qb.makeFilter(ctx, folderFilter)

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

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

	result, err := qb.queryGroupedFields(ctx, options, query)
	if err != nil {
		return nil, fmt.Errorf("error querying aggregate fields: %w", err)
	}

	idsResult, err := query.findIDs(ctx)
	if err != nil {
		return nil, fmt.Errorf("error finding IDs: %w", err)
	}

	result.IDs = make([]models.FolderID, len(idsResult))
	for i, id := range idsResult {
		result.IDs[i] = models.FolderID(id)
	}

	return result, nil
}

func (qb *FolderStore) queryGroupedFields(ctx context.Context, options models.FolderQueryOptions, query queryBuilder) (*models.FolderQueryResult, error) {
	if !options.Count {
		// nothing to do - return empty result
		return models.NewFolderQueryResult(qb), nil
	}

	aggregateQuery := qb.newQuery()

	if options.Count {
		aggregateQuery.addColumn("COUNT(DISTINCT temp.id) as total")
	}

	const includeSortPagination = false
	aggregateQuery.from = fmt.Sprintf("(%s) as temp", query.toSQL(includeSortPagination))

	out := struct {
		Total      int
		Duration   float64
		Megapixels float64
		Size       int64
	}{}
	if err := qb.repository.queryStruct(ctx, aggregateQuery.toSQL(includeSortPagination), query.args, &out); err != nil {
		return nil, err
	}

	ret := models.NewFolderQueryResult(qb)
	ret.Count = out.Total

	return ret, nil
}

var folderSortOptions = sortOptions{
	"created_at",
	"id",
	"path",
	"random",
	"updated_at",
}

func (qb *FolderStore) setQuerySort(query *queryBuilder, findFilter *models.FindFilterType) error {
	if findFilter == nil || findFilter.Sort == nil || *findFilter.Sort == "" {
		return nil
	}
	sort := findFilter.GetSort("path")

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

	direction := findFilter.GetDirection()
	query.sortAndPagination += getSort(sort, direction, "folders")

	return nil
}