stash/internal/autotag/studio.go

package autotag

import (
	"context"

	"github.com/stashapp/stash/pkg/gallery"
	"github.com/stashapp/stash/pkg/image"
	"github.com/stashapp/stash/pkg/match"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/scene"
	"github.com/stashapp/stash/pkg/txn"
)

// the following functions aren't used in Tagger because they assume
// use within a transaction

func addSceneStudio(ctx context.Context, sceneWriter scene.PartialUpdater, o *models.Scene, studioID int) (bool, error) {
	// don't set if already set
	if o.StudioID != nil {
		return false, nil
	}

	// set the studio id
	scenePartial := models.ScenePartial{
		StudioID: models.NewOptionalInt(studioID),
	}

	if _, err := sceneWriter.UpdatePartial(ctx, o.ID, scenePartial); err != nil {
		return false, err
	}
	return true, nil
}

func addImageStudio(ctx context.Context, imageWriter image.PartialUpdater, i *models.Image, studioID int) (bool, error) {
	// don't set if already set
	if i.StudioID != nil {
		return false, nil
	}

	// set the studio id
	imagePartial := models.ImagePartial{
		StudioID: models.NewOptionalInt(studioID),
	}

	if _, err := imageWriter.UpdatePartial(ctx, i.ID, imagePartial); err != nil {
		return false, err
	}
	return true, nil
}

func addGalleryStudio(ctx context.Context, galleryWriter GalleryFinderUpdater, o *models.Gallery, studioID int) (bool, error) {
	// don't set if already set
	if o.StudioID != nil {
		return false, nil
	}

	// set the studio id
	galleryPartial := models.GalleryPartial{
		StudioID: models.NewOptionalInt(studioID),
	}

	if _, err := galleryWriter.UpdatePartial(ctx, o.ID, galleryPartial); err != nil {
		return false, err
	}
	return true, nil
}

func getStudioTagger(p *models.Studio, aliases []string, cache *match.Cache) []tagger {
	ret := []tagger{{
		ID:    p.ID,
		Type:  "studio",
		Name:  p.Name.String,
		cache: cache,
	}}

	for _, a := range aliases {
		ret = append(ret, tagger{
			ID:   p.ID,
			Type: "studio",
			Name: a,
		})
	}

	return ret
}

type SceneFinderUpdater interface {
	scene.Queryer
	scene.PartialUpdater
}

// StudioScenes searches for scenes whose path matches the provided studio name and tags the scene with the studio, if studio is not already set on the scene.
func (tagger *Tagger) StudioScenes(ctx context.Context, p *models.Studio, paths []string, aliases []string, rw SceneFinderUpdater) error {
	t := getStudioTagger(p, aliases, tagger.Cache)

	for _, tt := range t {
		if err := tt.tagScenes(ctx, paths, rw, func(o *models.Scene) (bool, error) {
			// don't set if already set
			if o.StudioID != nil {
				return false, nil
			}

			// set the studio id
			scenePartial := models.ScenePartial{
				StudioID: models.NewOptionalInt(p.ID),
			}

			if err := txn.WithTxn(ctx, tagger.TxnManager, func(ctx context.Context) error {
				_, err := rw.UpdatePartial(ctx, o.ID, scenePartial)
				return err
			}); err != nil {
				return false, err
			}
			return true, nil
		}); err != nil {
			return err
		}
	}

	return nil
}

type ImageFinderUpdater interface {
	image.Queryer
	Find(ctx context.Context, id int) (*models.Image, error)
	UpdatePartial(ctx context.Context, id int, partial models.ImagePartial) (*models.Image, error)
}

// StudioImages searches for images whose path matches the provided studio name and tags the image with the studio, if studio is not already set on the image.
func (tagger *Tagger) StudioImages(ctx context.Context, p *models.Studio, paths []string, aliases []string, rw ImageFinderUpdater) error {
	t := getStudioTagger(p, aliases, tagger.Cache)

	for _, tt := range t {
		if err := tt.tagImages(ctx, paths, rw, func(i *models.Image) (bool, error) {
			// don't set if already set
			if i.StudioID != nil {
				return false, nil
			}

			// set the studio id
			imagePartial := models.ImagePartial{
				StudioID: models.NewOptionalInt(p.ID),
			}

			if err := txn.WithTxn(ctx, tagger.TxnManager, func(ctx context.Context) error {
				_, err := rw.UpdatePartial(ctx, i.ID, imagePartial)
				return err
			}); err != nil {
				return false, err
			}
			return true, nil
		}); err != nil {
			return err
		}
	}

	return nil
}

type GalleryFinderUpdater interface {
	gallery.Queryer
	gallery.PartialUpdater
	Find(ctx context.Context, id int) (*models.Gallery, error)
}

// StudioGalleries searches for galleries whose path matches the provided studio name and tags the gallery with the studio, if studio is not already set on the gallery.
func (tagger *Tagger) StudioGalleries(ctx context.Context, p *models.Studio, paths []string, aliases []string, rw GalleryFinderUpdater) error {
	t := getStudioTagger(p, aliases, tagger.Cache)

	for _, tt := range t {
		if err := tt.tagGalleries(ctx, paths, rw, func(o *models.Gallery) (bool, error) {
			// don't set if already set
			if o.StudioID != nil {
				return false, nil
			}

			// set the studio id
			galleryPartial := models.GalleryPartial{
				StudioID: models.NewOptionalInt(p.ID),
			}

			if err := txn.WithTxn(ctx, tagger.TxnManager, func(ctx context.Context) error {
				_, err := rw.UpdatePartial(ctx, o.ID, galleryPartial)
				return err
			}); err != nil {
				return false, err
			}
			return true, nil
		}); err != nil {
			return err
		}
	}

	return nil
}