package sqlx import ( "database/sql" "errors" "fmt" "time" "github.com/gofrs/uuid" "github.com/stashapp/stash-box/pkg/edit" "github.com/stashapp/stash-box/pkg/models" "github.com/stashapp/stash-box/pkg/utils" ) const ( studioTable = "studios" studioJoinKey = "studio_id" ) var ( studioDBTable = newTable(studioTable, func() interface{} { return &models.Studio{} }) studioURLTable = newTableJoin(studioTable, "studio_urls", studioJoinKey, func() interface{} { return &models.StudioURL{} }) studioRedirectTable = newTableJoin(studioTable, "studio_redirects", "source_id", func() interface{} { return &models.Redirect{} }) ) type studioQueryBuilder struct { dbi *dbi } func newStudioQueryBuilder(txn *txnState) models.StudioRepo { return &studioQueryBuilder{ dbi: newDBI(txn), } } func (qb *studioQueryBuilder) toModel(ro interface{}) *models.Studio { if ro != nil { return ro.(*models.Studio) } return nil } func (qb *studioQueryBuilder) Create(newStudio models.Studio) (*models.Studio, error) { ret, err := qb.dbi.Insert(studioDBTable, newStudio) return qb.toModel(ret), err } func (qb *studioQueryBuilder) Update(updatedStudio models.Studio) (*models.Studio, error) { ret, err := qb.dbi.Update(studioDBTable, updatedStudio, true) return qb.toModel(ret), err } func (qb *studioQueryBuilder) Destroy(id uuid.UUID) error { return qb.dbi.Delete(id, studioDBTable) } func (qb *studioQueryBuilder) CreateURLs(newJoins models.StudioURLs) error { return qb.dbi.InsertJoins(studioURLTable, &newJoins) } func (qb *studioQueryBuilder) CreateImages(newJoins models.StudiosImages) error { return qb.dbi.InsertJoins(studioImageTable, &newJoins) } func (qb *studioQueryBuilder) UpdateImages(studioID uuid.UUID, updatedJoins models.StudiosImages) error { return qb.dbi.ReplaceJoins(studioImageTable, studioID, &updatedJoins) } func (qb *studioQueryBuilder) UpdateURLs(studioID uuid.UUID, updatedJoins models.StudioURLs) error { return qb.dbi.ReplaceJoins(studioURLTable, studioID, &updatedJoins) } func (qb *studioQueryBuilder) Find(id uuid.UUID) (*models.Studio, error) { ret, err := qb.dbi.Find(id, studioDBTable) return qb.toModel(ret), err } func (qb *studioQueryBuilder) FindWithRedirect(id uuid.UUID) (*models.Studio, error) { query := ` SELECT S.* FROM studios S WHERE S.id = $1 AND S.deleted = FALSE UNION SELECT T.* FROM studio_redirects R JOIN studios T ON T.id = R.target_id WHERE R.source_id = $1 AND T.deleted = FALSE ` args := []interface{}{id} studios, err := qb.queryStudios(query, args) if len(studios) > 0 { return studios[0], err } return nil, err } func (qb *studioQueryBuilder) FindBySceneID(sceneID int) (models.Studios, error) { query := ` SELECT studios.* FROM studios LEFT JOIN scenes on scenes.studio_id = studios.id WHERE scenes.id = ? GROUP BY studios.id ` args := []interface{}{sceneID} return qb.queryStudios(query, args) } func (qb *studioQueryBuilder) FindByNames(names []string) (models.Studios, error) { query := "SELECT * FROM studios WHERE name IN " + getInBinding(len(names)) var args []interface{} for _, name := range names { args = append(args, name) } return qb.queryStudios(query, args) } func (qb *studioQueryBuilder) FindByName(name string) (*models.Studio, error) { query := "SELECT * FROM studios WHERE upper(name) = upper(?)" var args []interface{} args = append(args, name) results, err := qb.queryStudios(query, args) if err != nil || len(results) < 1 { return nil, err } return results[0], nil } func (qb *studioQueryBuilder) FindByParentID(id uuid.UUID) (models.Studios, error) { query := "SELECT * FROM studios WHERE parent_studio_id = ? AND deleted = FALSE" var args []interface{} args = append(args, id) return qb.queryStudios(query, args) } func (qb *studioQueryBuilder) Count() (int, error) { return runCountQuery(qb.dbi.db(), buildCountQuery("SELECT studios.id FROM studios"), nil) } func (qb *studioQueryBuilder) Query(filter models.StudioQueryInput, userID uuid.UUID) (models.Studios, int, error) { query := newQueryBuilder(studioDBTable) query.Body += "LEFT JOIN studios as parent_studio ON studios.parent_studio_id = parent_studio.id" query.Eq("studios.deleted", false) if q := filter.Name; q != nil && *q != "" { searchColumns := []string{"studios.name"} clause, thisArgs := getSearchBinding(searchColumns, *q, false, true) query.AddWhere(clause) query.AddArg(thisArgs...) } if q := filter.Names; q != nil && *q != "" { searchColumns := []string{"studios.name", "parent_studio.name"} clause, thisArgs := getSearchBinding(searchColumns, *q, false, true) query.AddWhere(clause) query.AddArg(thisArgs...) } if filter.HasParent != nil { if *filter.HasParent { query.AddWhere("parent_studio.id IS NOT NULL") } else { query.AddWhere("parent_studio.id IS NULL") } } if filter.IsFavorite != nil { // userID is internal based on user context so it is safe to append rather than bind q := fmt.Sprintf(" JOIN studio_favorites F ON studios.id = F.studio_id AND F.user_id = '%s'", userID) if *filter.IsFavorite { query.Body += q } else { query.Body += " LEFT" + q query.AddWhere("F.studio_id IS NULL") } } query.Sort = qb.getStudioSort(filter) query.Pagination = getPagination(filter.Page, filter.PerPage) var studios models.Studios countResult, err := qb.dbi.Query(*query, &studios) return studios, countResult, err } func (qb *studioQueryBuilder) getStudioSort(filter models.StudioQueryInput) string { secondary := "id" return getSort(filter.Sort.String(), filter.Direction.String(), "studios", &secondary) } func (qb *studioQueryBuilder) queryStudios(query string, args []interface{}) (models.Studios, error) { var output models.Studios err := qb.dbi.RawQuery(studioDBTable, query, args, &output) return output, err } func (qb *studioQueryBuilder) GetImages(id uuid.UUID) (models.StudiosImages, error) { joins := models.StudiosImages{} err := qb.dbi.FindJoins(studioImageTable, id, &joins) return joins, err } func (qb *studioQueryBuilder) GetURLs(id uuid.UUID) ([]*models.URL, error) { joins := models.StudioURLs{} err := qb.dbi.FindJoins(studioURLTable, id, &joins) urls := make([]*models.URL, len(joins)) for i, u := range joins { url := models.URL{ URL: u.URL, SiteID: u.SiteID, } urls[i] = &url } return urls, err } func (qb *studioQueryBuilder) GetAllURLs(ids []uuid.UUID) ([][]*models.URL, []error) { joins := models.StudioURLs{} err := qb.dbi.FindAllJoins(studioURLTable, ids, &joins) if err != nil { return nil, utils.DuplicateError(err, len(ids)) } m := make(map[uuid.UUID][]*models.URL) for _, join := range joins { url := models.URL{ URL: join.URL, SiteID: join.SiteID, } m[join.StudioID] = append(m[join.StudioID], &url) } result := make([][]*models.URL, len(ids)) for i, id := range ids { result[i] = m[id] } return result, nil } func (qb *studioQueryBuilder) CountByPerformer(performerID uuid.UUID) ([]*models.PerformerStudio, error) { var results []*models.PerformerStudio query := ` SELECT S.*, C.count FROM studios S JOIN ( SELECT studio_id, COUNT(*) FROM scene_performers SP JOIN scenes S ON SP.scene_id = S.id WHERE performer_id = ? GROUP BY studio_id ) C ON S.id = C.studio_id` query = qb.dbi.db().Rebind(query) if err := qb.dbi.db().Select(&results, query, performerID); err != nil && !errors.Is(err, sql.ErrNoRows) { return nil, err } return results, nil } func (qb *studioQueryBuilder) ApplyEdit(edit models.Edit, operation models.OperationEnum, studio *models.Studio) (*models.Studio, error) { data, err := edit.GetStudioData() if err != nil { return nil, err } switch operation { case models.OperationEnumCreate: now := time.Now() UUID, err := uuid.NewV4() if err != nil { return nil, err } newStudio := models.Studio{ ID: UUID, CreatedAt: now, UpdatedAt: now, } if data.New.Name == nil { return nil, errors.New("Missing studio name") } newStudio.CopyFromStudioEdit(*data.New, &models.StudioEdit{}) studio, err = qb.Create(newStudio) if err != nil { return nil, err } if len(data.New.AddedUrls) > 0 { urls := models.CreateStudioURLs(UUID, data.New.AddedUrls) if err := qb.CreateURLs(urls); err != nil { return nil, err } } if len(data.New.AddedImages) > 0 { images := models.CreateStudioImages(UUID, data.New.AddedImages) if err := qb.CreateImages(images); err != nil { return nil, err } } return studio, nil case models.OperationEnumDestroy: updatedStudio, err := qb.SoftDelete(*studio) if err != nil { return nil, err } if err = qb.deleteSceneStudios(studio.ID); err != nil { return nil, err } if err = qb.deleteStudioFavorites(studio.ID); err != nil { return nil, err } return updatedStudio, nil case models.OperationEnumModify: return qb.applyModifyEdit(studio, data) case models.OperationEnumMerge: updatedStudio, err := qb.applyModifyEdit(studio, data) if err != nil { return nil, err } for _, sourceID := range data.MergeSources { if err := qb.mergeInto(sourceID, studio.ID); err != nil { return nil, err } } return updatedStudio, nil default: return nil, errors.New("Unsupported operation: " + operation.String()) } } func (qb *studioQueryBuilder) applyModifyEdit(studio *models.Studio, data *models.StudioEditData) (*models.Studio, error) { if err := studio.ValidateModifyEdit(*data); err != nil { return nil, err } studio.CopyFromStudioEdit(*data.New, data.Old) updatedStudio, err := qb.Update(*studio) if err != nil { return nil, err } if err := qb.updateURLsFromEdit(updatedStudio, data); err != nil { return nil, err } if err := qb.updateImagesFromEdit(updatedStudio, data); err != nil { return nil, err } return updatedStudio, err } func (qb *studioQueryBuilder) GetEditURLs(id *uuid.UUID, data *models.StudioEdit) ([]*models.URL, error) { var urls []*models.URL if id != nil { currentURLs, err := qb.GetURLs(*id) if err != nil { return nil, err } urls = currentURLs } return edit.MergeURLs(urls, data.AddedUrls, data.RemovedUrls), nil } func (qb *studioQueryBuilder) updateURLsFromEdit(studio *models.Studio, data *models.StudioEditData) error { urls, err := qb.GetEditURLs(&studio.ID, data.New) if err != nil { return err } newURLs := models.CreateStudioURLs(studio.ID, urls) return qb.UpdateURLs(studio.ID, newURLs) } func (qb *studioQueryBuilder) GetEditImages(id *uuid.UUID, data *models.StudioEdit) ([]uuid.UUID, error) { var imageIds []uuid.UUID if id != nil { currentImages, err := qb.GetImages(*id) if err != nil { return nil, err } for _, v := range currentImages { imageIds = append(imageIds, v.ImageID) } } return utils.ProcessSlice(imageIds, data.AddedImages, data.RemovedImages), nil } func (qb *studioQueryBuilder) updateImagesFromEdit(studio *models.Studio, data *models.StudioEditData) error { ids, err := qb.GetEditImages(&studio.ID, data.New) if err != nil { return err } images := models.CreateStudioImages(studio.ID, ids) return qb.UpdateImages(studio.ID, images) } func (qb *studioQueryBuilder) mergeInto(sourceID uuid.UUID, targetID uuid.UUID) error { studio, err := qb.Find(sourceID) if err != nil { return err } if studio == nil { return errors.New("Merge source studio not found: " + sourceID.String()) } if studio.Deleted { return errors.New("Merge source studio is deleted: " + sourceID.String()) } _, err = qb.SoftDelete(*studio) if err != nil { return err } if err := qb.UpdateRedirects(sourceID, targetID); err != nil { return err } if err := qb.updateSceneStudios(sourceID, targetID); err != nil { return err } if err := qb.reassignFavorites(sourceID, targetID); err != nil { return err } redirect := models.Redirect{SourceID: sourceID, TargetID: targetID} return qb.CreateRedirect(redirect) } func (qb *studioQueryBuilder) CreateRedirect(newJoin models.Redirect) error { return qb.dbi.InsertJoin(studioRedirectTable, newJoin, nil) } func (qb *studioQueryBuilder) UpdateRedirects(oldTargetID uuid.UUID, newTargetID uuid.UUID) error { query := "UPDATE " + studioRedirectTable.table.Name() + " SET target_id = ? WHERE target_id = ?" args := []interface{}{newTargetID, oldTargetID} return qb.dbi.RawQuery(studioRedirectTable.table, query, args, nil) } func (qb *studioQueryBuilder) SoftDelete(studio models.Studio) (*models.Studio, error) { ret, err := qb.dbi.SoftDelete(studioDBTable, studio) return qb.toModel(ret), err } func (qb *studioQueryBuilder) updateSceneStudios(oldTargetID uuid.UUID, newTargetID uuid.UUID) error { // set existing studio ids to the new id query := `UPDATE ` + sceneDBTable.Name() + ` SET studio_id = ? WHERE studio_id = ?` args := []interface{}{newTargetID, oldTargetID} return qb.dbi.RawExec(query, args) } func (qb *studioQueryBuilder) reassignFavorites(oldTargetID uuid.UUID, newTargetID uuid.UUID) error { // Reassign performer favorites to new id where it isn't already assigned query := `UPDATE studio_favorites SET studio_id = ? WHERE studio_id = ? AND user_id NOT IN (SELECT user_id from studio_favorites WHERE studio_id = ?)` args := []interface{}{newTargetID, oldTargetID, newTargetID} err := qb.dbi.RawQuery(performerFavoriteTable.table, query, args, nil) if err != nil { return err } // Delete any remaining joins with the old performer query = `DELETE FROM studio_favorites WHERE studio_id = ?` args = []interface{}{oldTargetID} return qb.dbi.RawQuery(studioFavoriteTable.table, query, args, nil) } func (qb *studioQueryBuilder) deleteSceneStudios(id uuid.UUID) error { // set existing studio ids to null query := `UPDATE ` + sceneDBTable.Name() + ` SET studio_id = NULL WHERE studio_id = ?` args := []interface{}{id} return qb.dbi.RawExec(query, args) } func (qb *studioQueryBuilder) deleteStudioFavorites(id uuid.UUID) error { return qb.dbi.DeleteJoins(studioFavoriteTable, id) }