package api import ( "context" "errors" "fmt" "strconv" "github.com/stashapp/stash/pkg/match" "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/scraper" "github.com/stashapp/stash/pkg/sliceutil" "github.com/stashapp/stash/pkg/sliceutil/stringslice" ) func (r *queryResolver) ScrapeURL(ctx context.Context, url string, ty scraper.ScrapeContentType) (scraper.ScrapedContent, error) { return r.scraperCache().ScrapeURL(ctx, url, ty) } func (r *queryResolver) ListScrapers(ctx context.Context, types []scraper.ScrapeContentType) ([]*scraper.Scraper, error) { return r.scraperCache().ListScrapers(types), nil } func (r *queryResolver) ScrapePerformerURL(ctx context.Context, url string) (*models.ScrapedPerformer, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypePerformer) if err != nil { return nil, err } return marshalScrapedPerformer(content) } func (r *queryResolver) ScrapeSceneQuery(ctx context.Context, scraperID string, query string) ([]*models.ScrapedScene, error) { if query == "" { return nil, nil } content, err := r.scraperCache().ScrapeName(ctx, scraperID, query, scraper.ScrapeContentTypeScene) if err != nil { return nil, err } ret, err := marshalScrapedScenes(content) if err != nil { return nil, err } return ret, nil } func (r *queryResolver) ScrapeSceneURL(ctx context.Context, url string) (*models.ScrapedScene, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypeScene) if err != nil { return nil, err } ret, err := marshalScrapedScene(content) if err != nil { return nil, err } return ret, nil } func (r *queryResolver) ScrapeGalleryURL(ctx context.Context, url string) (*models.ScrapedGallery, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypeGallery) if err != nil { return nil, err } ret, err := marshalScrapedGallery(content) if err != nil { return nil, err } return ret, nil } func (r *queryResolver) ScrapeImageURL(ctx context.Context, url string) (*models.ScrapedImage, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypeImage) if err != nil { return nil, err } return marshalScrapedImage(content) } func (r *queryResolver) ScrapeMovieURL(ctx context.Context, url string) (*models.ScrapedMovie, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypeMovie) if err != nil { return nil, err } ret, err := marshalScrapedMovie(content) if err != nil { return nil, err } return ret, nil } func (r *queryResolver) ScrapeGroupURL(ctx context.Context, url string) (*models.ScrapedGroup, error) { content, err := r.scraperCache().ScrapeURL(ctx, url, scraper.ScrapeContentTypeMovie) if err != nil { return nil, err } ret, err := marshalScrapedMovie(content) if err != nil { return nil, err } // convert to scraped group group := &models.ScrapedGroup{ StoredID: ret.StoredID, Name: ret.Name, Aliases: ret.Aliases, Duration: ret.Duration, Date: ret.Date, Rating: ret.Rating, Director: ret.Director, URLs: ret.URLs, Synopsis: ret.Synopsis, Studio: ret.Studio, Tags: ret.Tags, FrontImage: ret.FrontImage, BackImage: ret.BackImage, } return group, nil } func (r *queryResolver) ScrapeSingleScene(ctx context.Context, source scraper.Source, input ScrapeSingleSceneInput) ([]*models.ScrapedScene, error) { var ret []*models.ScrapedScene var sceneID int if input.SceneID != nil { var err error sceneID, err = strconv.Atoi(*input.SceneID) if err != nil { return nil, fmt.Errorf("%w: sceneID is not an integer: '%s'", ErrInput, *input.SceneID) } } switch { case source.ScraperID != nil: var err error var c scraper.ScrapedContent var content []scraper.ScrapedContent switch { case input.SceneID != nil: c, err = r.scraperCache().ScrapeID(ctx, *source.ScraperID, sceneID, scraper.ScrapeContentTypeScene) if c != nil { content = []scraper.ScrapedContent{c} } case input.SceneInput != nil: c, err = r.scraperCache().ScrapeFragment(ctx, *source.ScraperID, scraper.Input{Scene: input.SceneInput}) if c != nil { content = []scraper.ScrapedContent{c} } case input.Query != nil: content, err = r.scraperCache().ScrapeName(ctx, *source.ScraperID, *input.Query, scraper.ScrapeContentTypeScene) default: err = fmt.Errorf("%w: scene_id, scene_input, or query must be set", ErrInput) } if err != nil { return nil, err } ret, err = marshalScrapedScenes(content) if err != nil { return nil, err } case source.StashBoxIndex != nil || source.StashBoxEndpoint != nil: b, err := resolveStashBox(source.StashBoxIndex, source.StashBoxEndpoint) if err != nil { return nil, err } client := r.newStashBoxClient(*b) switch { case input.SceneID != nil: var fps []models.Fingerprints fps, err = r.getScenesFingerprints(ctx, []int{sceneID}) if err != nil { return nil, err } ret, err = client.FindSceneByFingerprints(ctx, fps[0]) case input.Query != nil: ret, err = client.QueryScene(ctx, *input.Query) default: return nil, fmt.Errorf("%w: scene_id or query must be set", ErrInput) } if err != nil { return nil, err } // TODO - this should happen after any scene is scraped if err := r.matchScenesRelationships(ctx, ret, *source.StashBoxEndpoint); err != nil { return nil, err } default: return nil, fmt.Errorf("%w: scraper_id or stash_box_index must be set", ErrInput) } return ret, nil } func (r *queryResolver) ScrapeMultiScenes(ctx context.Context, source scraper.Source, input ScrapeMultiScenesInput) ([][]*models.ScrapedScene, error) { if source.ScraperID != nil { return nil, ErrNotImplemented } else if source.StashBoxIndex != nil || source.StashBoxEndpoint != nil { b, err := resolveStashBox(source.StashBoxIndex, source.StashBoxEndpoint) if err != nil { return nil, err } client := r.newStashBoxClient(*b) sceneIDs, err := stringslice.StringSliceToIntSlice(input.SceneIds) if err != nil { return nil, err } fps, err := r.getScenesFingerprints(ctx, sceneIDs) if err != nil { return nil, err } ret, err := client.FindScenesByFingerprints(ctx, fps) if err != nil { return nil, err } // match relationships - this mutates the existing scenes so we can // just flatten the slice and pass it in flat := sliceutil.Flatten(ret) if err := r.matchScenesRelationships(ctx, flat, *source.StashBoxEndpoint); err != nil { return nil, err } return ret, nil } return nil, errors.New("scraper_id or stash_box_index must be set") } func (r *queryResolver) getScenesFingerprints(ctx context.Context, ids []int) ([]models.Fingerprints, error) { fingerprints := make([]models.Fingerprints, len(ids)) if err := r.withReadTxn(ctx, func(ctx context.Context) error { qb := r.repository.Scene for i, sceneID := range ids { scene, err := qb.Find(ctx, sceneID) if err != nil { return err } if scene == nil { return fmt.Errorf("scene with id %d not found", sceneID) } if err := scene.LoadFiles(ctx, qb); err != nil { return err } var sceneFPs models.Fingerprints for _, f := range scene.Files.List() { sceneFPs = append(sceneFPs, f.Fingerprints...) } fingerprints[i] = sceneFPs } return nil }); err != nil { return nil, err } return fingerprints, nil } // matchSceneRelationships accepts scraped scenes and attempts to match its relationships to existing stash models. func (r *queryResolver) matchScenesRelationships(ctx context.Context, ss []*models.ScrapedScene, endpoint string) error { if err := r.withReadTxn(ctx, func(ctx context.Context) error { matcher := match.SceneRelationships{ PerformerFinder: r.repository.Performer, TagFinder: r.repository.Tag, StudioFinder: r.repository.Studio, } for _, s := range ss { if err := matcher.MatchRelationships(ctx, s, endpoint); err != nil { return err } } return nil }); err != nil { return err } return nil } func (r *queryResolver) ScrapeSingleStudio(ctx context.Context, source scraper.Source, input ScrapeSingleStudioInput) ([]*models.ScrapedStudio, error) { if source.StashBoxIndex != nil || source.StashBoxEndpoint != nil { b, err := resolveStashBox(source.StashBoxIndex, source.StashBoxEndpoint) if err != nil { return nil, err } client := r.newStashBoxClient(*b) var ret []*models.ScrapedStudio out, err := client.FindStudio(ctx, *input.Query) if err != nil { return nil, err } else if out != nil { ret = append(ret, out) } if len(ret) > 0 { if err := r.withReadTxn(ctx, func(ctx context.Context) error { for _, studio := range ret { if err := match.ScrapedStudioHierarchy(ctx, r.repository.Studio, studio, *source.StashBoxEndpoint); err != nil { return err } } return nil }); err != nil { return nil, err } return ret, nil } return nil, nil } return nil, errors.New("stash_box_index must be set") } func (r *queryResolver) ScrapeSinglePerformer(ctx context.Context, source scraper.Source, input ScrapeSinglePerformerInput) ([]*models.ScrapedPerformer, error) { var ret []*models.ScrapedPerformer switch { case source.ScraperID != nil: switch { case input.PerformerInput != nil: performer, err := r.scraperCache().ScrapeFragment(ctx, *source.ScraperID, scraper.Input{Performer: input.PerformerInput}) if err != nil { return nil, err } ret, err = marshalScrapedPerformers([]scraper.ScrapedContent{performer}) if err != nil { return nil, err } case input.Query != nil: content, err := r.scraperCache().ScrapeName(ctx, *source.ScraperID, *input.Query, scraper.ScrapeContentTypePerformer) if err != nil { return nil, err } ret, err = marshalScrapedPerformers(content) if err != nil { return nil, err } default: return nil, ErrNotImplemented } case source.StashBoxIndex != nil || source.StashBoxEndpoint != nil: b, err := resolveStashBox(source.StashBoxIndex, source.StashBoxEndpoint) if err != nil { return nil, err } client := r.newStashBoxClient(*b) var query string switch { case input.PerformerID != nil: names, err := r.findPerformerNames(ctx, []string{*input.PerformerID}) if err != nil { return nil, err } query = names[0] case input.Query != nil: query = *input.Query default: return nil, ErrNotImplemented } if query == "" { return nil, nil } ret, err = client.QueryPerformer(ctx, query) if err != nil { return nil, err } default: return nil, errors.New("scraper_id or stash_box_index must be set") } return ret, nil } func (r *queryResolver) ScrapeMultiPerformers(ctx context.Context, source scraper.Source, input ScrapeMultiPerformersInput) ([][]*models.ScrapedPerformer, error) { if source.ScraperID != nil { return nil, ErrNotImplemented } else if source.StashBoxIndex != nil || source.StashBoxEndpoint != nil { names, err := r.findPerformerNames(ctx, input.PerformerIds) if err != nil { return nil, err } b, err := resolveStashBox(source.StashBoxIndex, source.StashBoxEndpoint) if err != nil { return nil, err } client := r.newStashBoxClient(*b) return client.QueryPerformers(ctx, names) } return nil, errors.New("scraper_id or stash_box_index must be set") } func (r *queryResolver) findPerformerNames(ctx context.Context, performerIDs []string) ([]string, error) { ids, err := stringslice.StringSliceToIntSlice(performerIDs) if err != nil { return nil, err } names := make([]string, len(ids)) if err := r.withReadTxn(ctx, func(ctx context.Context) error { p, err := r.repository.Performer.FindMany(ctx, ids) if err != nil { return err } for i, pp := range p { names[i] = pp.Name } return nil }); err != nil { return nil, err } return names, nil } func (r *queryResolver) ScrapeSingleGallery(ctx context.Context, source scraper.Source, input ScrapeSingleGalleryInput) ([]*models.ScrapedGallery, error) { var ret []*models.ScrapedGallery if source.StashBoxIndex != nil || source.StashBoxEndpoint != nil { return nil, ErrNotSupported } if source.ScraperID == nil { return nil, fmt.Errorf("%w: scraper_id must be set", ErrInput) } var c scraper.ScrapedContent switch { case input.GalleryID != nil: galleryID, err := strconv.Atoi(*input.GalleryID) if err != nil { return nil, fmt.Errorf("%w: gallery id is not an integer: '%s'", ErrInput, *input.GalleryID) } c, err = r.scraperCache().ScrapeID(ctx, *source.ScraperID, galleryID, scraper.ScrapeContentTypeGallery) if err != nil { return nil, err } ret, err = marshalScrapedGalleries([]scraper.ScrapedContent{c}) if err != nil { return nil, err } case input.GalleryInput != nil: c, err := r.scraperCache().ScrapeFragment(ctx, *source.ScraperID, scraper.Input{Gallery: input.GalleryInput}) if err != nil { return nil, err } ret, err = marshalScrapedGalleries([]scraper.ScrapedContent{c}) if err != nil { return nil, err } default: return nil, ErrNotImplemented } return ret, nil } func (r *queryResolver) ScrapeSingleImage(ctx context.Context, source scraper.Source, input ScrapeSingleImageInput) ([]*models.ScrapedImage, error) { if source.StashBoxIndex != nil { return nil, ErrNotSupported } if source.ScraperID == nil { return nil, fmt.Errorf("%w: scraper_id must be set", ErrInput) } var c scraper.ScrapedContent switch { case input.ImageID != nil: imageID, err := strconv.Atoi(*input.ImageID) if err != nil { return nil, fmt.Errorf("%w: image id is not an integer: '%s'", ErrInput, *input.ImageID) } c, err = r.scraperCache().ScrapeID(ctx, *source.ScraperID, imageID, scraper.ScrapeContentTypeImage) if err != nil { return nil, err } return marshalScrapedImages([]scraper.ScrapedContent{c}) case input.ImageInput != nil: c, err := r.scraperCache().ScrapeFragment(ctx, *source.ScraperID, scraper.Input{Image: input.ImageInput}) if err != nil { return nil, err } return marshalScrapedImages([]scraper.ScrapedContent{c}) default: return nil, ErrNotImplemented } } func (r *queryResolver) ScrapeSingleMovie(ctx context.Context, source scraper.Source, input ScrapeSingleMovieInput) ([]*models.ScrapedMovie, error) { return nil, ErrNotSupported } func (r *queryResolver) ScrapeSingleGroup(ctx context.Context, source scraper.Source, input ScrapeSingleGroupInput) ([]*models.ScrapedGroup, error) { return nil, ErrNotSupported }