stash/pkg/scraper/stash.go

package scraper

import (
	"context"
	"errors"
	"strconv"

	"github.com/jinzhu/copier"
	"github.com/shurcooL/graphql"

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

type stashScraper struct {
	scraper      scraperTypeConfig
	config       config
	globalConfig GlobalConfig
	txnManager   models.TransactionManager
}

func newStashScraper(scraper scraperTypeConfig, txnManager models.TransactionManager, config config, globalConfig GlobalConfig) *stashScraper {
	return &stashScraper{
		scraper:      scraper,
		config:       config,
		globalConfig: globalConfig,
		txnManager:   txnManager,
	}
}

func (s *stashScraper) getStashClient() *graphql.Client {
	url := s.config.StashServer.URL
	return graphql.NewClient(url+"/graphql", nil)
}

type stashFindPerformerNamePerformer struct {
	ID   string `json:"id" graphql:"id"`
	Name string `json:"name" graphql:"name"`
}

func (p stashFindPerformerNamePerformer) toPerformer() *models.ScrapedPerformer {
	return &models.ScrapedPerformer{
		Name: &p.Name,
		// put id into the URL field
		URL: &p.ID,
	}
}

type stashFindPerformerNamesResultType struct {
	Count      int                                `graphql:"count"`
	Performers []*stashFindPerformerNamePerformer `graphql:"performers"`
}

func (s *stashScraper) scrapePerformersByName(name string) ([]*models.ScrapedPerformer, error) {
	client := s.getStashClient()

	var q struct {
		FindPerformers stashFindPerformerNamesResultType `graphql:"findPerformers(filter: $f)"`
	}

	page := 1
	perPage := 10

	vars := map[string]interface{}{
		"f": models.FindFilterType{
			Q:       &name,
			Page:    &page,
			PerPage: &perPage,
		},
	}

	err := client.Query(context.Background(), &q, vars)
	if err != nil {
		return nil, err
	}

	var ret []*models.ScrapedPerformer
	for _, p := range q.FindPerformers.Performers {
		ret = append(ret, p.toPerformer())
	}

	return ret, nil
}

func (s *stashScraper) scrapePerformerByFragment(scrapedPerformer models.ScrapedPerformerInput) (*models.ScrapedPerformer, error) {
	client := s.getStashClient()

	var q struct {
		FindPerformer *models.ScrapedPerformerStash `graphql:"findPerformer(id: $f)"`
	}

	performerID := *scrapedPerformer.URL

	// get the id from the URL field
	vars := map[string]interface{}{
		"f": performerID,
	}

	err := client.Query(context.Background(), &q, vars)
	if err != nil {
		return nil, err
	}

	if q.FindPerformer != nil {
		// the ids of the tags must be nilled
		for _, t := range q.FindPerformer.Tags {
			t.ID = nil
		}
	}

	// need to copy back to a scraped performer
	ret := models.ScrapedPerformer{}
	err = copier.Copy(&ret, q.FindPerformer)
	if err != nil {
		return nil, err
	}

	// get the performer image directly
	ret.Image, err = getStashPerformerImage(s.config.StashServer.URL, performerID, s.globalConfig)
	if err != nil {
		return nil, err
	}

	return &ret, nil
}

func (s *stashScraper) scrapeSceneByFragment(scene models.SceneUpdateInput) (*models.ScrapedScene, error) {
	// query by MD5
	// assumes that the scene exists in the database
	id, err := strconv.Atoi(scene.ID)
	if err != nil {
		return nil, err
	}

	var storedScene *models.Scene
	if err := s.txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		var err error
		storedScene, err = r.Scene().Find(id)
		return err
	}); err != nil {
		return nil, err
	}

	var q struct {
		FindScene *models.ScrapedSceneStash `graphql:"findSceneByHash(input: $c)"`
	}

	type SceneHashInput struct {
		Checksum *string `graphql:"checksum" json:"checksum"`
		Oshash   *string `graphql:"oshash" json:"oshash"`
	}

	input := SceneHashInput{
		Checksum: &storedScene.Checksum.String,
		Oshash:   &storedScene.OSHash.String,
	}

	vars := map[string]interface{}{
		"c": &input,
	}

	client := s.getStashClient()
	err = client.Query(context.Background(), &q, vars)
	if err != nil {
		return nil, err
	}

	if q.FindScene != nil {
		// the ids of the studio, performers and tags must be nilled
		if q.FindScene.Studio != nil {
			q.FindScene.Studio.ID = nil
		}

		for _, p := range q.FindScene.Performers {
			p.ID = nil
		}

		for _, t := range q.FindScene.Tags {
			t.ID = nil
		}
	}

	// need to copy back to a scraped scene
	ret := models.ScrapedScene{}
	err = copier.Copy(&ret, q.FindScene)
	if err != nil {
		return nil, err
	}

	// get the performer image directly
	ret.Image, err = getStashSceneImage(s.config.StashServer.URL, q.FindScene.ID, s.globalConfig)
	if err != nil {
		return nil, err
	}

	return &ret, nil
}

func (s *stashScraper) scrapeGalleryByFragment(scene models.GalleryUpdateInput) (*models.ScrapedGallery, error) {
	id, err := strconv.Atoi(scene.ID)
	if err != nil {
		return nil, err
	}

	// query by MD5
	// assumes that the gallery exists in the database
	var storedGallery *models.Gallery
	if err := s.txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		qb := r.Gallery()

		var err error
		storedGallery, err = qb.Find(id)
		return err
	}); err != nil {
		return nil, err
	}

	var q struct {
		FindGallery *models.ScrapedGalleryStash `graphql:"findGalleryByHash(input: $c)"`
	}

	type GalleryHashInput struct {
		Checksum *string `graphql:"checksum" json:"checksum"`
	}

	input := GalleryHashInput{
		Checksum: &storedGallery.Checksum,
	}

	vars := map[string]interface{}{
		"c": &input,
	}

	client := s.getStashClient()
	err = client.Query(context.Background(), &q, vars)
	if err != nil {
		return nil, err
	}

	if q.FindGallery != nil {
		// the ids of the studio, performers and tags must be nilled
		if q.FindGallery.Studio != nil {
			q.FindGallery.Studio.ID = nil
		}

		for _, p := range q.FindGallery.Performers {
			p.ID = nil
		}

		for _, t := range q.FindGallery.Tags {
			t.ID = nil
		}
	}

	// need to copy back to a scraped scene
	ret := models.ScrapedGallery{}
	err = copier.Copy(&ret, q.FindGallery)
	if err != nil {
		return nil, err
	}

	return &ret, nil
}

func (s *stashScraper) scrapePerformerByURL(url string) (*models.ScrapedPerformer, error) {
	return nil, errors.New("scrapePerformerByURL not supported for stash scraper")
}

func (s *stashScraper) scrapeSceneByURL(url string) (*models.ScrapedScene, error) {
	return nil, errors.New("scrapeSceneByURL not supported for stash scraper")
}

func (s *stashScraper) scrapeGalleryByURL(url string) (*models.ScrapedGallery, error) {
	return nil, errors.New("scrapeGalleryByURL not supported for stash scraper")
}

func (s *stashScraper) scrapeMovieByURL(url string) (*models.ScrapedMovie, error) {
	return nil, errors.New("scrapeMovieByURL not supported for stash scraper")
}

func sceneFromUpdateFragment(scene models.SceneUpdateInput, txnManager models.TransactionManager) (*models.Scene, error) {
	id, err := strconv.Atoi(scene.ID)
	if err != nil {
		return nil, err
	}

	// TODO - should we modify it with the input?
	var ret *models.Scene
	if err := txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		var err error
		ret, err = r.Scene().Find(id)
		return err
	}); err != nil {
		return nil, err
	}
	return ret, nil
}

func galleryFromUpdateFragment(gallery models.GalleryUpdateInput, txnManager models.TransactionManager) (ret *models.Gallery, err error) {
	id, err := strconv.Atoi(gallery.ID)
	if err != nil {
		return nil, err
	}

	if err := txnManager.WithReadTxn(context.TODO(), func(r models.ReaderRepository) error {
		ret, err = r.Gallery().Find(id)
		return err
	}); err != nil {
		return nil, err
	}

	return ret, nil
}