Skip to content

Commit

Permalink
Add stash scraper type (stashapp#269)
Browse files Browse the repository at this point in the history
* Add stash scraper type

* Add graphql client to vendor

* Embed stash credentials in URL

* Fill URL from scraped scene

* Nil IDs returned from remote stash

* Nil check
  • Loading branch information
WithoutPants authored and Leopere committed Dec 21, 2019
1 parent 0870d9b commit 282dfc5
Show file tree
Hide file tree
Showing 21 changed files with 1,526 additions and 52 deletions.
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ require (
github.com/inconshreveable/mousetrap v1.0.0 // indirect
github.com/jmoiron/sqlx v1.2.0
github.com/mattn/go-sqlite3 v1.10.0
github.com/microcosm-cc/bluemonday v1.0.2 // indirect
github.com/rs/cors v1.6.0
github.com/shurcooL/graphql v0.0.0-20181231061246-d48a9a75455f
github.com/sirupsen/logrus v1.4.2
github.com/spf13/pflag v1.0.3
github.com/spf13/viper v1.4.0
Expand Down
50 changes: 4 additions & 46 deletions go.sum

Large diffs are not rendered by default.

12 changes: 12 additions & 0 deletions gqlgen.yml
Original file line number Diff line number Diff line change
Expand Up @@ -28,3 +28,15 @@ models:
model: github.com/stashapp/stash/pkg/models.Studio
Tag:
model: github.com/stashapp/stash/pkg/models.Tag
ScrapedPerformer:
model: github.com/stashapp/stash/pkg/models.ScrapedPerformer
ScrapedScene:
model: github.com/stashapp/stash/pkg/models.ScrapedScene
ScrapedScenePerformer:
model: github.com/stashapp/stash/pkg/models.ScrapedScenePerformer
ScrapedSceneStudio:
model: github.com/stashapp/stash/pkg/models.ScrapedSceneStudio
ScrapedSceneTag:
model: github.com/stashapp/stash/pkg/models.ScrapedSceneTag
SceneFileType:
model: github.com/stashapp/stash/pkg/models.SceneFileType
2 changes: 1 addition & 1 deletion pkg/api/resolver_query_find_scene.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,10 @@ import (

func (r *queryResolver) FindScene(ctx context.Context, id *string, checksum *string) (*models.Scene, error) {
qb := models.NewSceneQueryBuilder()
idInt, _ := strconv.Atoi(*id)
var scene *models.Scene
var err error
if id != nil {
idInt, _ := strconv.Atoi(*id)
scene, err = qb.Find(idInt)
} else if checksum != nil {
scene, err = qb.FindByChecksum(*checksum)
Expand Down
11 changes: 11 additions & 0 deletions pkg/models/model_scene.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,14 @@ func (s Scene) GetTitle() string {

return filepath.Base(s.Path)
}

type SceneFileType struct {
Size *string `graphql:"size" json:"size"`
Duration *float64 `graphql:"duration" json:"duration"`
VideoCodec *string `graphql:"video_codec" json:"video_codec"`
AudioCodec *string `graphql:"audio_codec" json:"audio_codec"`
Width *int `graphql:"width" json:"width"`
Height *int `graphql:"height" json:"height"`
Framerate *float64 `graphql:"framerate" json:"framerate"`
Bitrate *int `graphql:"bitrate" json:"bitrate"`
}
62 changes: 62 additions & 0 deletions pkg/models/model_scraped_item.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,65 @@ type ScrapedItem struct {
CreatedAt SQLiteTimestamp `db:"created_at" json:"created_at"`
UpdatedAt SQLiteTimestamp `db:"updated_at" json:"updated_at"`
}

type ScrapedPerformer struct {
Name *string `graphql:"name" json:"name"`
URL *string `graphql:"url" json:"url"`
Twitter *string `graphql:"twitter" json:"twitter"`
Instagram *string `graphql:"instagram" json:"instagram"`
Birthdate *string `graphql:"birthdate" json:"birthdate"`
Ethnicity *string `graphql:"ethnicity" json:"ethnicity"`
Country *string `graphql:"country" json:"country"`
EyeColor *string `graphql:"eye_color" json:"eye_color"`
Height *string `graphql:"height" json:"height"`
Measurements *string `graphql:"measurements" json:"measurements"`
FakeTits *string `graphql:"fake_tits" json:"fake_tits"`
CareerLength *string `graphql:"career_length" json:"career_length"`
Tattoos *string `graphql:"tattoos" json:"tattoos"`
Piercings *string `graphql:"piercings" json:"piercings"`
Aliases *string `graphql:"aliases" json:"aliases"`
}

type ScrapedScene struct {
Title *string `graphql:"title" json:"title"`
Details *string `graphql:"details" json:"details"`
URL *string `graphql:"url" json:"url"`
Date *string `graphql:"date" json:"date"`
File *SceneFileType `graphql:"file" json:"file"`
Studio *ScrapedSceneStudio `graphql:"studio" json:"studio"`
Tags []*ScrapedSceneTag `graphql:"tags" json:"tags"`
Performers []*ScrapedScenePerformer `graphql:"performers" json:"performers"`
}

type ScrapedScenePerformer struct {
// Set if performer matched
ID *string `graphql:"id" json:"id"`
Name string `graphql:"name" json:"name"`
URL *string `graphql:"url" json:"url"`
Twitter *string `graphql:"twitter" json:"twitter"`
Instagram *string `graphql:"instagram" json:"instagram"`
Birthdate *string `graphql:"birthdate" json:"birthdate"`
Ethnicity *string `graphql:"ethnicity" json:"ethnicity"`
Country *string `graphql:"country" json:"country"`
EyeColor *string `graphql:"eye_color" json:"eye_color"`
Height *string `graphql:"height" json:"height"`
Measurements *string `graphql:"measurements" json:"measurements"`
FakeTits *string `graphql:"fake_tits" json:"fake_tits"`
CareerLength *string `graphql:"career_length" json:"career_length"`
Tattoos *string `graphql:"tattoos" json:"tattoos"`
Piercings *string `graphql:"piercings" json:"piercings"`
Aliases *string `graphql:"aliases" json:"aliases"`
}

type ScrapedSceneStudio struct {
// Set if studio matched
ID *string `graphql:"id" json:"id"`
Name string `graphql:"name" json:"name"`
URL *string `graphql:"url" json:"url"`
}

type ScrapedSceneTag struct {
// Set if tag matched
ID *string `graphql:"id" json:"id"`
Name string `graphql:"name" json:"name"`
}
20 changes: 20 additions & 0 deletions pkg/scraper/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,20 @@ import (
"github.com/stashapp/stash/pkg/models"
)

type stashServer struct {
URL string `yaml:"url"`
}

type scraperAction string

const (
scraperActionScript scraperAction = "script"
scraperActionStash scraperAction = "stash"
)

var allScraperAction = []scraperAction{
scraperActionScript,
scraperActionStash,
}

func (e scraperAction) IsValid() bool {
Expand All @@ -31,6 +37,8 @@ func (e scraperAction) IsValid() bool {
type scraperTypeConfig struct {
Action scraperAction `yaml:"action"`
Script []string `yaml:"script,flow"`

scraperConfig *scraperConfig
}

type scrapePerformerNamesFunc func(c scraperTypeConfig, name string) ([]*models.ScrapedPerformer, error)
Expand All @@ -43,6 +51,8 @@ type performerByNameConfig struct {
func (c *performerByNameConfig) resolveFn() {
if c.Action == scraperActionScript {
c.performScrape = scrapePerformerNamesScript
} else if c.Action == scraperActionStash {
c.performScrape = scrapePerformerNamesStash
}
}

Expand All @@ -56,6 +66,8 @@ type performerByFragmentConfig struct {
func (c *performerByFragmentConfig) resolveFn() {
if c.Action == scraperActionScript {
c.performScrape = scrapePerformerFragmentScript
} else if c.Action == scraperActionStash {
c.performScrape = scrapePerformerFragmentStash
}
}

Expand Down Expand Up @@ -97,6 +109,8 @@ type sceneByFragmentConfig struct {
func (c *sceneByFragmentConfig) resolveFn() {
if c.Action == scraperActionScript {
c.performScrape = scrapeSceneFragmentScript
} else if c.Action == scraperActionStash {
c.performScrape = scrapeSceneFragmentStash
}
}

Expand All @@ -121,6 +135,7 @@ type scraperConfig struct {
PerformerByURL []*scrapePerformerByURLConfig `yaml:"performerByURL"`
SceneByFragment *sceneByFragmentConfig `yaml:"sceneByFragment"`
SceneByURL []*scrapeSceneByURLConfig `yaml:"sceneByURL"`
StashServer *stashServer `yaml:"stashServer"`
}

func loadScraperFromYAML(path string) (*scraperConfig, error) {
Expand Down Expand Up @@ -152,19 +167,24 @@ func loadScraperFromYAML(path string) (*scraperConfig, error) {
func (c *scraperConfig) initialiseConfigs() {
if c.PerformerByName != nil {
c.PerformerByName.resolveFn()
c.PerformerByName.scraperConfig = c
}
if c.PerformerByFragment != nil {
c.PerformerByFragment.resolveFn()
c.PerformerByFragment.scraperConfig = c
}
for _, s := range c.PerformerByURL {
s.resolveFn()
s.scraperConfig = c
}

if c.SceneByFragment != nil {
c.SceneByFragment.resolveFn()
c.SceneByFragment.scraperConfig = c
}
for _, s := range c.SceneByURL {
s.resolveFn()
s.scraperConfig = c
}
}

Expand Down
8 changes: 5 additions & 3 deletions pkg/scraper/scrapers.go
Original file line number Diff line number Diff line change
Expand Up @@ -216,9 +216,11 @@ func ScrapeScene(scraperID string, scene models.SceneUpdateInput) (*models.Scrap
return nil, err
}

err = postScrapeScene(ret)
if err != nil {
return nil, err
if ret != nil {
err = postScrapeScene(ret)
if err != nil {
return nil, err
}
}

return ret, nil
Expand Down
132 changes: 132 additions & 0 deletions pkg/scraper/stash.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
package scraper

import (
"context"
"strconv"

"github.com/shurcooL/graphql"

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

func getStashClient(c scraperTypeConfig) *graphql.Client {
url := c.scraperConfig.StashServer.URL
return graphql.NewClient(url+"/graphql", nil)
}

type stashFindPerformerNamePerformer struct {
ID string `json:"id" graphql:"id"`
Name string `json:"id" 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 scrapePerformerNamesStash(c scraperTypeConfig, name string) ([]*models.ScrapedPerformer, error) {
client := getStashClient(c)

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 scrapePerformerFragmentStash(c scraperTypeConfig, scrapedPerformer models.ScrapedPerformerInput) (*models.ScrapedPerformer, error) {
client := getStashClient(c)

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

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

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

return q.FindPerformer, nil
}

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

storedScene, err := qb.Find(id)

if err != nil {
return nil, err
}

var q struct {
FindScene *models.ScrapedScene `graphql:"findScene(checksum: $c)"`
}

checksum := graphql.String(storedScene.Checksum)
vars := map[string]interface{}{
"c": &checksum,
}

client := getStashClient(c)
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
}
}

return q.FindScene, nil
}
4 changes: 4 additions & 0 deletions ui/v2/src/components/scenes/SceneDetails/SceneEditPanel.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,10 @@ export const SceneEditPanel: FunctionComponent<IProps> = (props: IProps) => {
if (!date && scene.date) {
setDate(scene.date);
}

if (!url && scene.url) {
setUrl(scene.url);
}

if (!studioId && scene.studio && scene.studio.id) {
setStudioId(scene.studio.id);
Expand Down
16 changes: 16 additions & 0 deletions vendor/github.com/shurcooL/graphql/.travis.yml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit 282dfc5

Please sign in to comment.