mirror of https://github.com/stashapp/stash.git
133 lines
3.3 KiB
Go
133 lines
3.3 KiB
Go
package autotag
|
|
|
|
import (
|
|
"database/sql"
|
|
|
|
"github.com/stashapp/stash/pkg/models"
|
|
)
|
|
|
|
func getMatchingStudios(path string, reader models.StudioReader) ([]*models.Studio, error) {
|
|
words := getPathWords(path)
|
|
candidates, err := reader.QueryForAutoTag(words)
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var ret []*models.Studio
|
|
for _, c := range candidates {
|
|
if nameMatchesPath(c.Name.String, path) {
|
|
ret = append(ret, c)
|
|
}
|
|
}
|
|
|
|
return ret, nil
|
|
}
|
|
|
|
func addSceneStudio(sceneWriter models.SceneReaderWriter, sceneID, studioID int) (bool, error) {
|
|
// don't set if already set
|
|
scene, err := sceneWriter.Find(sceneID)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
if scene.StudioID.Valid {
|
|
return false, nil
|
|
}
|
|
|
|
// set the studio id
|
|
s := sql.NullInt64{Int64: int64(studioID), Valid: true}
|
|
scenePartial := models.ScenePartial{
|
|
ID: sceneID,
|
|
StudioID: &s,
|
|
}
|
|
|
|
if _, err := sceneWriter.Update(scenePartial); err != nil {
|
|
return false, err
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
func addImageStudio(imageWriter models.ImageReaderWriter, imageID, studioID int) (bool, error) {
|
|
// don't set if already set
|
|
image, err := imageWriter.Find(imageID)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
if image.StudioID.Valid {
|
|
return false, nil
|
|
}
|
|
|
|
// set the studio id
|
|
s := sql.NullInt64{Int64: int64(studioID), Valid: true}
|
|
imagePartial := models.ImagePartial{
|
|
ID: imageID,
|
|
StudioID: &s,
|
|
}
|
|
|
|
if _, err := imageWriter.Update(imagePartial); err != nil {
|
|
return false, err
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
func addGalleryStudio(galleryWriter models.GalleryReaderWriter, galleryID, studioID int) (bool, error) {
|
|
// don't set if already set
|
|
gallery, err := galleryWriter.Find(galleryID)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
if gallery.StudioID.Valid {
|
|
return false, nil
|
|
}
|
|
|
|
// set the studio id
|
|
s := sql.NullInt64{Int64: int64(studioID), Valid: true}
|
|
galleryPartial := models.GalleryPartial{
|
|
ID: galleryID,
|
|
StudioID: &s,
|
|
}
|
|
|
|
if _, err := galleryWriter.UpdatePartial(galleryPartial); err != nil {
|
|
return false, err
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
func getStudioTagger(p *models.Studio) tagger {
|
|
return tagger{
|
|
ID: p.ID,
|
|
Type: "studio",
|
|
Name: p.Name.String,
|
|
}
|
|
}
|
|
|
|
// 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 StudioScenes(p *models.Studio, paths []string, rw models.SceneReaderWriter) error {
|
|
t := getStudioTagger(p)
|
|
|
|
return t.tagScenes(paths, rw, func(subjectID, otherID int) (bool, error) {
|
|
return addSceneStudio(rw, otherID, subjectID)
|
|
})
|
|
}
|
|
|
|
// 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 StudioImages(p *models.Studio, paths []string, rw models.ImageReaderWriter) error {
|
|
t := getStudioTagger(p)
|
|
|
|
return t.tagImages(paths, rw, func(subjectID, otherID int) (bool, error) {
|
|
return addImageStudio(rw, otherID, subjectID)
|
|
})
|
|
}
|
|
|
|
// 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 StudioGalleries(p *models.Studio, paths []string, rw models.GalleryReaderWriter) error {
|
|
t := getStudioTagger(p)
|
|
|
|
return t.tagGalleries(paths, rw, func(subjectID, otherID int) (bool, error) {
|
|
return addGalleryStudio(rw, otherID, subjectID)
|
|
})
|
|
}
|