stash/pkg/match/scraped.go

package match

import (
	"context"
	"strconv"

	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/performer"
	"github.com/stashapp/stash/pkg/studio"
	"github.com/stashapp/stash/pkg/tag"
)

type PerformerFinder interface {
	models.PerformerQueryer
	FindByNames(ctx context.Context, names []string, nocase bool) ([]*models.Performer, error)
	FindByStashID(ctx context.Context, stashID models.StashID) ([]*models.Performer, error)
}

type GroupNamesFinder interface {
	FindByNames(ctx context.Context, names []string, nocase bool) ([]*models.Group, error)
}

type SceneRelationships struct {
	PerformerFinder PerformerFinder
	TagFinder       models.TagQueryer
	StudioFinder    StudioFinder
}

// MatchRelationships accepts a scraped scene and attempts to match its relationships to existing stash models.
func (r SceneRelationships) MatchRelationships(ctx context.Context, s *models.ScrapedScene, endpoint string) error {
	thisStudio := s.Studio
	for thisStudio != nil {
		if err := ScrapedStudio(ctx, r.StudioFinder, thisStudio, endpoint); err != nil {
			return err
		}

		thisStudio = thisStudio.Parent
	}

	for _, p := range s.Performers {
		err := ScrapedPerformer(ctx, r.PerformerFinder, p, endpoint)
		if err != nil {
			return err
		}
	}

	for _, t := range s.Tags {
		err := ScrapedTag(ctx, r.TagFinder, t, endpoint)
		if err != nil {
			return err
		}
	}

	return nil
}

// ScrapedPerformer matches the provided performer with the
// performers in the database and sets the ID field if one is found.
func ScrapedPerformer(ctx context.Context, qb PerformerFinder, p *models.ScrapedPerformer, stashBoxEndpoint string) error {
	if p.StoredID != nil || p.Name == nil {
		return nil
	}

	// Check if a performer with the StashID already exists
	if stashBoxEndpoint != "" && p.RemoteSiteID != nil {
		performers, err := qb.FindByStashID(ctx, models.StashID{
			StashID:  *p.RemoteSiteID,
			Endpoint: stashBoxEndpoint,
		})
		if err != nil {
			return err
		}
		if len(performers) > 0 {
			id := strconv.Itoa(performers[0].ID)
			p.StoredID = &id
			return nil
		}
	}

	performers, err := qb.FindByNames(ctx, []string{*p.Name}, true)
	if err != nil {
		return err
	}

	if len(performers) == 0 {
		// if no names matched, try match an exact alias
		performers, err = performer.ByAlias(ctx, qb, *p.Name)
		if err != nil {
			return err
		}
	}

	if len(performers) != 1 {
		// ignore - cannot match
		return nil
	}

	id := strconv.Itoa(performers[0].ID)
	p.StoredID = &id
	return nil
}

type StudioFinder interface {
	models.StudioQueryer
	FindByStashID(ctx context.Context, stashID models.StashID) ([]*models.Studio, error)
}

// ScrapedStudio matches the provided studio with the studios
// in the database and sets the ID field if one is found.
func ScrapedStudio(ctx context.Context, qb StudioFinder, s *models.ScrapedStudio, stashBoxEndpoint string) error {
	if s.StoredID != nil {
		return nil
	}

	// Check if a studio with the StashID already exists
	if stashBoxEndpoint != "" && s.RemoteSiteID != nil {
		studios, err := qb.FindByStashID(ctx, models.StashID{
			StashID:  *s.RemoteSiteID,
			Endpoint: stashBoxEndpoint,
		})
		if err != nil {
			return err
		}
		if len(studios) > 0 {
			id := strconv.Itoa(studios[0].ID)
			s.StoredID = &id
			return nil
		}
	}

	st, err := studio.ByName(ctx, qb, s.Name)

	if err != nil {
		return err
	}

	if st == nil {
		// try matching by alias
		st, err = studio.ByAlias(ctx, qb, s.Name)
		if err != nil {
			return err
		}
	}

	if st == nil {
		// ignore - cannot match
		return nil
	}

	id := strconv.Itoa(st.ID)
	s.StoredID = &id
	return nil
}

// ScrapedStudioHierarchy executes ScrapedStudio for the provided studio and its parents recursively.
func ScrapedStudioHierarchy(ctx context.Context, qb StudioFinder, s *models.ScrapedStudio, stashBoxEndpoint string) error {
	if err := ScrapedStudio(ctx, qb, s, stashBoxEndpoint); err != nil {
		return err
	}

	if s.Parent == nil {
		return nil
	}

	return ScrapedStudioHierarchy(ctx, qb, s.Parent, stashBoxEndpoint)
}

// ScrapedGroup matches the provided movie with the movies
// in the database and returns the ID field if one is found.
func ScrapedGroup(ctx context.Context, qb GroupNamesFinder, storedID *string, name *string) (matchedID *string, err error) {
	if storedID != nil || name == nil {
		return
	}

	movies, err := qb.FindByNames(ctx, []string{*name}, true)

	if err != nil {
		return
	}

	if len(movies) != 1 {
		// ignore - cannot match
		return
	}

	id := strconv.Itoa(movies[0].ID)
	matchedID = &id
	return
}

// ScrapedTag matches the provided tag with the tags
// in the database and sets the ID field if one is found.
func ScrapedTag(ctx context.Context, qb models.TagQueryer, s *models.ScrapedTag, stashBoxEndpoint string) error {
	if s.StoredID != nil {
		return nil
	}

	// Check if a tag with the StashID already exists
	if stashBoxEndpoint != "" && s.RemoteSiteID != nil {
		if finder, ok := qb.(models.TagFinder); ok {
			tags, err := finder.FindByStashID(ctx, models.StashID{
				StashID:  *s.RemoteSiteID,
				Endpoint: stashBoxEndpoint,
			})
			if err != nil {
				return err
			}
			if len(tags) > 0 {
				id := strconv.Itoa(tags[0].ID)
				s.StoredID = &id
				return nil
			}
		}
	}

	t, err := tag.ByName(ctx, qb, s.Name)

	if err != nil {
		return err
	}

	if t == nil {
		// try matching by alias
		t, err = tag.ByAlias(ctx, qb, s.Name)
		if err != nil {
			return err
		}
	}

	if t == nil {
		// ignore - cannot match
		return nil
	}

	id := strconv.Itoa(t.ID)
	s.StoredID = &id
	return nil
}