package manager import ( "context" "fmt" "math" "runtime" "strconv" "sync" "time" "github.com/jmoiron/sqlx" "github.com/stashapp/stash/pkg/database" "github.com/stashapp/stash/pkg/logger" "github.com/stashapp/stash/pkg/manager/jsonschema" "github.com/stashapp/stash/pkg/manager/paths" "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/utils" ) type ExportTask struct { Mappings *jsonschema.Mappings Scraped []jsonschema.ScrapedItem } func (t *ExportTask) Start(wg *sync.WaitGroup) { defer wg.Done() // @manager.total = Scene.count + Gallery.count + Performer.count + Studio.count + Movie.count workerCount := runtime.GOMAXPROCS(0) // set worker count to number of cpus available t.Mappings = &jsonschema.Mappings{} t.Scraped = []jsonschema.ScrapedItem{} ctx := context.TODO() startTime := time.Now() paths.EnsureJSONDirs() t.ExportScenes(ctx, workerCount) t.ExportGalleries(ctx) t.ExportPerformers(ctx, workerCount) t.ExportStudios(ctx, workerCount) t.ExportMovies(ctx, workerCount) if err := instance.JSON.saveMappings(t.Mappings); err != nil { logger.Errorf("[mappings] failed to save json: %s", err.Error()) } t.ExportScrapedItems(ctx) logger.Infof("Export complete in %s.", time.Since(startTime)) } func (t *ExportTask) ExportScenes(ctx context.Context, workers int) { var scenesWg sync.WaitGroup qb := models.NewSceneQueryBuilder() scenes, err := qb.All() if err != nil { logger.Errorf("[scenes] failed to fetch all scenes: %s", err.Error()) } jobCh := make(chan *models.Scene, workers*2) // make a buffered channel to feed workers logger.Info("[scenes] exporting") startTime := time.Now() for w := 0; w < workers; w++ { // create export Scene workers scenesWg.Add(1) go exportScene(&scenesWg, jobCh, t, nil) // no db data is changed so tx is set to nil } for i, scene := range scenes { index := i + 1 if (i % 100) == 0 { // make progress easier to read logger.Progressf("[scenes] %d of %d", index, len(scenes)) } t.Mappings.Scenes = append(t.Mappings.Scenes, jsonschema.PathMapping{Path: scene.Path, Checksum: scene.Checksum}) jobCh <- scene // feed workers } close(jobCh) // close channel so that workers will know no more jobs are available scenesWg.Wait() logger.Infof("[scenes] export complete in %s. %d workers used.", time.Since(startTime), workers) } func exportScene(wg *sync.WaitGroup, jobChan <-chan *models.Scene, t *ExportTask, tx *sqlx.Tx) { defer wg.Done() sceneQB := models.NewSceneQueryBuilder() studioQB := models.NewStudioQueryBuilder() movieQB := models.NewMovieQueryBuilder() galleryQB := models.NewGalleryQueryBuilder() performerQB := models.NewPerformerQueryBuilder() tagQB := models.NewTagQueryBuilder() sceneMarkerQB := models.NewSceneMarkerQueryBuilder() joinQB := models.NewJoinsQueryBuilder() for scene := range jobChan { newSceneJSON := jsonschema.Scene{ CreatedAt: models.JSONTime{Time: scene.CreatedAt.Timestamp}, UpdatedAt: models.JSONTime{Time: scene.UpdatedAt.Timestamp}, } var studioName string if scene.StudioID.Valid { studio, _ := studioQB.Find(int(scene.StudioID.Int64), tx) if studio != nil { studioName = studio.Name.String } } var galleryChecksum string gallery, _ := galleryQB.FindBySceneID(scene.ID, tx) if gallery != nil { galleryChecksum = gallery.Checksum } performers, _ := performerQB.FindNameBySceneID(scene.ID, tx) sceneMovies, _ := joinQB.GetSceneMovies(scene.ID, tx) tags, _ := tagQB.FindBySceneID(scene.ID, tx) sceneMarkers, _ := sceneMarkerQB.FindBySceneID(scene.ID, tx) if scene.Title.Valid { newSceneJSON.Title = scene.Title.String } if studioName != "" { newSceneJSON.Studio = studioName } if scene.URL.Valid { newSceneJSON.URL = scene.URL.String } if scene.Date.Valid { newSceneJSON.Date = utils.GetYMDFromDatabaseDate(scene.Date.String) } if scene.Rating.Valid { newSceneJSON.Rating = int(scene.Rating.Int64) } newSceneJSON.OCounter = scene.OCounter if scene.Details.Valid { newSceneJSON.Details = scene.Details.String } if galleryChecksum != "" { newSceneJSON.Gallery = galleryChecksum } newSceneJSON.Performers = t.getPerformerNames(performers) newSceneJSON.Tags = t.getTagNames(tags) for _, sceneMarker := range sceneMarkers { primaryTag, err := tagQB.Find(sceneMarker.PrimaryTagID, tx) if err != nil { logger.Errorf("[scenes] <%s> invalid primary tag for scene marker: %s", scene.Checksum, err.Error()) continue } sceneMarkerTags, err := tagQB.FindBySceneMarkerID(sceneMarker.ID, tx) if err != nil { logger.Errorf("[scenes] <%s> invalid tags for scene marker: %s", scene.Checksum, err.Error()) continue } if sceneMarker.Title == "" || sceneMarker.Seconds == 0 || primaryTag.Name == "" { logger.Errorf("[scenes] invalid scene marker: %v", sceneMarker) } sceneMarkerJSON := jsonschema.SceneMarker{ Title: sceneMarker.Title, Seconds: t.getDecimalString(sceneMarker.Seconds), PrimaryTag: primaryTag.Name, Tags: t.getTagNames(sceneMarkerTags), CreatedAt: models.JSONTime{Time: sceneMarker.CreatedAt.Timestamp}, UpdatedAt: models.JSONTime{Time: sceneMarker.UpdatedAt.Timestamp}, } newSceneJSON.Markers = append(newSceneJSON.Markers, sceneMarkerJSON) } for _, sceneMovie := range sceneMovies { movie, _ := movieQB.Find(sceneMovie.MovieID, tx) if movie.Name.Valid { sceneMovieJSON := jsonschema.SceneMovie{ MovieName: movie.Name.String, SceneIndex: int(sceneMovie.SceneIndex.Int64), } newSceneJSON.Movies = append(newSceneJSON.Movies, sceneMovieJSON) } } newSceneJSON.File = &jsonschema.SceneFile{} if scene.Size.Valid { newSceneJSON.File.Size = scene.Size.String } if scene.Duration.Valid { newSceneJSON.File.Duration = t.getDecimalString(scene.Duration.Float64) } if scene.VideoCodec.Valid { newSceneJSON.File.VideoCodec = scene.VideoCodec.String } if scene.AudioCodec.Valid { newSceneJSON.File.AudioCodec = scene.AudioCodec.String } if scene.Format.Valid { newSceneJSON.File.Format = scene.Format.String } if scene.Width.Valid { newSceneJSON.File.Width = int(scene.Width.Int64) } if scene.Height.Valid { newSceneJSON.File.Height = int(scene.Height.Int64) } if scene.Framerate.Valid { newSceneJSON.File.Framerate = t.getDecimalString(scene.Framerate.Float64) } if scene.Bitrate.Valid { newSceneJSON.File.Bitrate = int(scene.Bitrate.Int64) } cover, err := sceneQB.GetSceneCover(scene.ID, tx) if err != nil { logger.Errorf("[scenes] <%s> error getting scene cover: %s", scene.Checksum, err.Error()) continue } if len(cover) > 0 { newSceneJSON.Cover = utils.GetBase64StringFromData(cover) } sceneJSON, err := instance.JSON.getScene(scene.Checksum) if err != nil { logger.Debugf("[scenes] error reading scene json: %s", err.Error()) } else if jsonschema.CompareJSON(*sceneJSON, newSceneJSON) { continue } if err := instance.JSON.saveScene(scene.Checksum, &newSceneJSON); err != nil { logger.Errorf("[scenes] <%s> failed to save json: %s", scene.Checksum, err.Error()) } } } func (t *ExportTask) ExportGalleries(ctx context.Context) { qb := models.NewGalleryQueryBuilder() galleries, err := qb.All() if err != nil { logger.Errorf("[galleries] failed to fetch all galleries: %s", err.Error()) } logger.Info("[galleries] exporting") for i, gallery := range galleries { index := i + 1 logger.Progressf("[galleries] %d of %d", index, len(galleries)) t.Mappings.Galleries = append(t.Mappings.Galleries, jsonschema.PathMapping{Path: gallery.Path, Checksum: gallery.Checksum}) } logger.Infof("[galleries] export complete") } func (t *ExportTask) ExportPerformers(ctx context.Context, workers int) { var performersWg sync.WaitGroup qb := models.NewPerformerQueryBuilder() performers, err := qb.All() if err != nil { logger.Errorf("[performers] failed to fetch all performers: %s", err.Error()) } jobCh := make(chan *models.Performer, workers*2) // make a buffered channel to feed workers logger.Info("[performers] exporting") startTime := time.Now() for w := 0; w < workers; w++ { // create export Performer workers performersWg.Add(1) go exportPerformer(&performersWg, jobCh) } for i, performer := range performers { index := i + 1 logger.Progressf("[performers] %d of %d", index, len(performers)) t.Mappings.Performers = append(t.Mappings.Performers, jsonschema.NameMapping{Name: performer.Name.String, Checksum: performer.Checksum}) jobCh <- performer // feed workers } close(jobCh) // close channel so workers will know that no more jobs are available performersWg.Wait() logger.Infof("[performers] export complete in %s. %d workers used.", time.Since(startTime), workers) } func exportPerformer(wg *sync.WaitGroup, jobChan <-chan *models.Performer) { defer wg.Done() performerQB := models.NewPerformerQueryBuilder() for performer := range jobChan { newPerformerJSON := jsonschema.Performer{ CreatedAt: models.JSONTime{Time: performer.CreatedAt.Timestamp}, UpdatedAt: models.JSONTime{Time: performer.UpdatedAt.Timestamp}, } if performer.Name.Valid { newPerformerJSON.Name = performer.Name.String } if performer.Gender.Valid { newPerformerJSON.Gender = performer.Gender.String } if performer.URL.Valid { newPerformerJSON.URL = performer.URL.String } if performer.Birthdate.Valid { newPerformerJSON.Birthdate = utils.GetYMDFromDatabaseDate(performer.Birthdate.String) } if performer.Ethnicity.Valid { newPerformerJSON.Ethnicity = performer.Ethnicity.String } if performer.Country.Valid { newPerformerJSON.Country = performer.Country.String } if performer.EyeColor.Valid { newPerformerJSON.EyeColor = performer.EyeColor.String } if performer.Height.Valid { newPerformerJSON.Height = performer.Height.String } if performer.Measurements.Valid { newPerformerJSON.Measurements = performer.Measurements.String } if performer.FakeTits.Valid { newPerformerJSON.FakeTits = performer.FakeTits.String } if performer.CareerLength.Valid { newPerformerJSON.CareerLength = performer.CareerLength.String } if performer.Tattoos.Valid { newPerformerJSON.Tattoos = performer.Tattoos.String } if performer.Piercings.Valid { newPerformerJSON.Piercings = performer.Piercings.String } if performer.Aliases.Valid { newPerformerJSON.Aliases = performer.Aliases.String } if performer.Twitter.Valid { newPerformerJSON.Twitter = performer.Twitter.String } if performer.Instagram.Valid { newPerformerJSON.Instagram = performer.Instagram.String } if performer.Favorite.Valid { newPerformerJSON.Favorite = performer.Favorite.Bool } image, err := performerQB.GetPerformerImage(performer.ID, nil) if err != nil { logger.Errorf("[performers] <%s> error getting performers image: %s", performer.Checksum, err.Error()) continue } if len(image) > 0 { newPerformerJSON.Image = utils.GetBase64StringFromData(image) } performerJSON, err := instance.JSON.getPerformer(performer.Checksum) if err != nil { logger.Debugf("[performers] error reading performer json: %s", err.Error()) } else if jsonschema.CompareJSON(*performerJSON, newPerformerJSON) { continue } if err := instance.JSON.savePerformer(performer.Checksum, &newPerformerJSON); err != nil { logger.Errorf("[performers] <%s> failed to save json: %s", performer.Checksum, err.Error()) } } } func (t *ExportTask) ExportStudios(ctx context.Context, workers int) { var studiosWg sync.WaitGroup qb := models.NewStudioQueryBuilder() studios, err := qb.All() if err != nil { logger.Errorf("[studios] failed to fetch all studios: %s", err.Error()) } logger.Info("[studios] exporting") startTime := time.Now() jobCh := make(chan *models.Studio, workers*2) // make a buffered channel to feed workers for w := 0; w < workers; w++ { // create export Studio workers studiosWg.Add(1) go exportStudio(&studiosWg, jobCh) } for i, studio := range studios { index := i + 1 logger.Progressf("[studios] %d of %d", index, len(studios)) t.Mappings.Studios = append(t.Mappings.Studios, jsonschema.NameMapping{Name: studio.Name.String, Checksum: studio.Checksum}) jobCh <- studio // feed workers } close(jobCh) studiosWg.Wait() logger.Infof("[studios] export complete in %s. %d workers used.", time.Since(startTime), workers) } func exportStudio(wg *sync.WaitGroup, jobChan <-chan *models.Studio) { defer wg.Done() studioQB := models.NewStudioQueryBuilder() for studio := range jobChan { newStudioJSON := jsonschema.Studio{ CreatedAt: models.JSONTime{Time: studio.CreatedAt.Timestamp}, UpdatedAt: models.JSONTime{Time: studio.UpdatedAt.Timestamp}, } if studio.Name.Valid { newStudioJSON.Name = studio.Name.String } if studio.URL.Valid { newStudioJSON.URL = studio.URL.String } if studio.ParentID.Valid { parent, _ := studioQB.Find(int(studio.ParentID.Int64), nil) if parent != nil { newStudioJSON.ParentStudio = parent.Name.String } } image, err := studioQB.GetStudioImage(studio.ID, nil) if err != nil { logger.Errorf("[studios] <%s> error getting studio image: %s", studio.Checksum, err.Error()) continue } if len(image) > 0 { newStudioJSON.Image = utils.GetBase64StringFromData(image) } studioJSON, err := instance.JSON.getStudio(studio.Checksum) if err != nil { logger.Debugf("[studios] error reading studio json: %s", err.Error()) } else if jsonschema.CompareJSON(*studioJSON, newStudioJSON) { continue } if err := instance.JSON.saveStudio(studio.Checksum, &newStudioJSON); err != nil { logger.Errorf("[studios] <%s> failed to save json: %s", studio.Checksum, err.Error()) } } } func (t *ExportTask) ExportMovies(ctx context.Context, workers int) { var moviesWg sync.WaitGroup qb := models.NewMovieQueryBuilder() movies, err := qb.All() if err != nil { logger.Errorf("[movies] failed to fetch all movies: %s", err.Error()) } logger.Info("[movies] exporting") startTime := time.Now() jobCh := make(chan *models.Movie, workers*2) // make a buffered channel to feed workers for w := 0; w < workers; w++ { // create export Studio workers moviesWg.Add(1) go exportMovie(&moviesWg, jobCh) } for i, movie := range movies { index := i + 1 logger.Progressf("[movies] %d of %d", index, len(movies)) t.Mappings.Movies = append(t.Mappings.Movies, jsonschema.NameMapping{Name: movie.Name.String, Checksum: movie.Checksum}) jobCh <- movie // feed workers } close(jobCh) moviesWg.Wait() logger.Infof("[movies] export complete in %s. %d workers used.", time.Since(startTime), workers) } func exportMovie(wg *sync.WaitGroup, jobChan <-chan *models.Movie) { defer wg.Done() movieQB := models.NewMovieQueryBuilder() studioQB := models.NewStudioQueryBuilder() for movie := range jobChan { newMovieJSON := jsonschema.Movie{ CreatedAt: models.JSONTime{Time: movie.CreatedAt.Timestamp}, UpdatedAt: models.JSONTime{Time: movie.UpdatedAt.Timestamp}, } if movie.Name.Valid { newMovieJSON.Name = movie.Name.String } if movie.Aliases.Valid { newMovieJSON.Aliases = movie.Aliases.String } if movie.Date.Valid { newMovieJSON.Date = utils.GetYMDFromDatabaseDate(movie.Date.String) } if movie.Rating.Valid { newMovieJSON.Rating = int(movie.Rating.Int64) } if movie.Duration.Valid { newMovieJSON.Duration = int(movie.Duration.Int64) } if movie.Director.Valid { newMovieJSON.Director = movie.Director.String } if movie.Synopsis.Valid { newMovieJSON.Synopsis = movie.Synopsis.String } if movie.URL.Valid { newMovieJSON.URL = movie.URL.String } if movie.StudioID.Valid { studio, _ := studioQB.Find(int(movie.StudioID.Int64), nil) if studio != nil { newMovieJSON.Studio = studio.Name.String } } frontImage, err := movieQB.GetFrontImage(movie.ID, nil) if err != nil { logger.Errorf("[movies] <%s> error getting movie front image: %s", movie.Checksum, err.Error()) continue } if len(frontImage) > 0 { newMovieJSON.FrontImage = utils.GetBase64StringFromData(frontImage) } backImage, err := movieQB.GetBackImage(movie.ID, nil) if err != nil { logger.Errorf("[movies] <%s> error getting movie back image: %s", movie.Checksum, err.Error()) continue } if len(backImage) > 0 { newMovieJSON.BackImage = utils.GetBase64StringFromData(backImage) } movieJSON, err := instance.JSON.getMovie(movie.Checksum) if err != nil { logger.Debugf("[movies] error reading movie json: %s", err.Error()) } else if jsonschema.CompareJSON(*movieJSON, newMovieJSON) { continue } if err := instance.JSON.saveMovie(movie.Checksum, &newMovieJSON); err != nil { logger.Errorf("[movies] <%s> failed to save json: %s", movie.Checksum, err.Error()) } } } func (t *ExportTask) ExportScrapedItems(ctx context.Context) { tx := database.DB.MustBeginTx(ctx, nil) defer tx.Commit() qb := models.NewScrapedItemQueryBuilder() sqb := models.NewStudioQueryBuilder() scrapedItems, err := qb.All() if err != nil { logger.Errorf("[scraped sites] failed to fetch all items: %s", err.Error()) } logger.Info("[scraped sites] exporting") for i, scrapedItem := range scrapedItems { index := i + 1 logger.Progressf("[scraped sites] %d of %d", index, len(scrapedItems)) var studioName string if scrapedItem.StudioID.Valid { studio, _ := sqb.Find(int(scrapedItem.StudioID.Int64), tx) if studio != nil { studioName = studio.Name.String } } newScrapedItemJSON := jsonschema.ScrapedItem{} if scrapedItem.Title.Valid { newScrapedItemJSON.Title = scrapedItem.Title.String } if scrapedItem.Description.Valid { newScrapedItemJSON.Description = scrapedItem.Description.String } if scrapedItem.URL.Valid { newScrapedItemJSON.URL = scrapedItem.URL.String } if scrapedItem.Date.Valid { newScrapedItemJSON.Date = utils.GetYMDFromDatabaseDate(scrapedItem.Date.String) } if scrapedItem.Rating.Valid { newScrapedItemJSON.Rating = scrapedItem.Rating.String } if scrapedItem.Tags.Valid { newScrapedItemJSON.Tags = scrapedItem.Tags.String } if scrapedItem.Models.Valid { newScrapedItemJSON.Models = scrapedItem.Models.String } if scrapedItem.Episode.Valid { newScrapedItemJSON.Episode = int(scrapedItem.Episode.Int64) } if scrapedItem.GalleryFilename.Valid { newScrapedItemJSON.GalleryFilename = scrapedItem.GalleryFilename.String } if scrapedItem.GalleryURL.Valid { newScrapedItemJSON.GalleryURL = scrapedItem.GalleryURL.String } if scrapedItem.VideoFilename.Valid { newScrapedItemJSON.VideoFilename = scrapedItem.VideoFilename.String } if scrapedItem.VideoURL.Valid { newScrapedItemJSON.VideoURL = scrapedItem.VideoURL.String } newScrapedItemJSON.Studio = studioName updatedAt := models.JSONTime{Time: scrapedItem.UpdatedAt.Timestamp} // TODO keeping ruby format newScrapedItemJSON.UpdatedAt = updatedAt t.Scraped = append(t.Scraped, newScrapedItemJSON) } scrapedJSON, err := instance.JSON.getScraped() if err != nil { logger.Debugf("[scraped sites] error reading json: %s", err.Error()) } if !jsonschema.CompareJSON(scrapedJSON, t.Scraped) { if err := instance.JSON.saveScaped(t.Scraped); err != nil { logger.Errorf("[scraped sites] failed to save json: %s", err.Error()) } } logger.Infof("[scraped sites] export complete") } func (t *ExportTask) getPerformerNames(performers []*models.Performer) []string { if len(performers) == 0 { return nil } var results []string for _, performer := range performers { if performer.Name.Valid { results = append(results, performer.Name.String) } } return results } func (t *ExportTask) getTagNames(tags []*models.Tag) []string { if len(tags) == 0 { return nil } var results []string for _, tag := range tags { if tag.Name != "" { results = append(results, tag.Name) } } return results } func (t *ExportTask) getDecimalString(num float64) string { if num == 0 { return "" } precision := getPrecision(num) if precision == 0 { precision = 1 } return fmt.Sprintf("%."+strconv.Itoa(precision)+"f", num) } func getPrecision(num float64) int { if num == 0 { return 0 } e := 1.0 p := 0 for (math.Round(num*e) / e) != num { e *= 10 p++ } return p }