Add tags to studios (#4858)

* Fix makeTagFilter mode

* Remove studio_tags filter criterion

This is handled by studios_filter. The support for this still needs to be added in the UI, so I have removed the criterion options in the short-term.
---------
Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>

pkg/sqlite/database.go

@@ -30,7 +30,7 @@ const (
 	dbConnTimeout = 30
 )
 
-var appSchemaVersion uint = 62
+var appSchemaVersion uint = 63
 
 //go:embed migrations/*.sql
 var migrationsBox embed.FS

pkg/sqlite/migrations/63_studio_tags.up.sql

@@ -0,0 +1,9 @@
+CREATE TABLE `studios_tags` (
+  `studio_id` integer NOT NULL,
+  `tag_id` integer NOT NULL,
+  foreign key(`studio_id`) references `studios`(`id`) on delete CASCADE,
+  foreign key(`tag_id`) references `tags`(`id`) on delete CASCADE,
+  PRIMARY KEY(`studio_id`, `tag_id`)
+);
+
+CREATE INDEX `index_studios_tags_on_tag_id` on `studios_tags` (`tag_id`);

pkg/sqlite/setup_test.go

@@ -207,6 +207,9 @@ const (
 	tagIdxWithPerformer
 	tagIdx1WithPerformer
 	tagIdx2WithPerformer
+	tagIdxWithStudio
+	tagIdx1WithStudio
+	tagIdx2WithStudio
 	tagIdxWithGallery
 	tagIdx1WithGallery
 	tagIdx2WithGallery
@@ -245,6 +248,10 @@ const (
 	studioIdxWithScenePerformer
 	studioIdxWithImagePerformer
 	studioIdxWithGalleryPerformer
+	studioIdxWithTag
+	studioIdx2WithTag
+	studioIdxWithTwoTags
+	studioIdxWithParentTag
 	studioIdxWithGrandChild
 	studioIdxWithParentAndChild
 	studioIdxWithGrandParent
@@ -510,6 +517,15 @@ var (
 	}
 )
 
+var (
+	studioTags = linkMap{
+		studioIdxWithTag:       {tagIdxWithStudio},
+		studioIdx2WithTag:      {tagIdx2WithStudio},
+		studioIdxWithTwoTags:   {tagIdx1WithStudio, tagIdx2WithStudio},
+		studioIdxWithParentTag: {tagIdxWithParentAndChild},
+	}
+)
+
 var (
 	performerTags = linkMap{
 		performerIdxWithTag:       {tagIdxWithPerformer},
@@ -1566,6 +1582,11 @@ func getTagPerformerCount(id int) int {
 	return len(performerTags.reverseLookup(idx))
 }
 
+func getTagStudioCount(id int) int {
+	idx := indexFromID(tagIDs, id)
+	return len(studioTags.reverseLookup(idx))
+}
+
 func getTagParentCount(id int) int {
 	if id == tagIDs[tagIdxWithParentTag] || id == tagIDs[tagIdxWithGrandParent] || id == tagIDs[tagIdxWithParentAndChild] {
 		return 1
@@ -1681,11 +1702,13 @@ func createStudios(ctx context.Context, n int, o int) error {
 		// studios [ i ] and [ n + o - i - 1  ] should have similar names with only the Name!=NaMe part different
 
 		name = getStudioStringValue(index, name)
+		tids := indexesToIDs(tagIDs, studioTags[i])
 		studio := models.Studio{
 			Name:          name,
 			URL:           getStudioStringValue(index, urlField),
 			Favorite:      getStudioBoolValue(index),
 			IgnoreAutoTag: getIgnoreAutoTag(i),
+			TagIDs:        models.NewRelatedIDs(tids),
 		}
 		// only add aliases for some scenes
 		if i == studioIdxWithMovie || i%5 == 0 {

pkg/sqlite/studio.go

@@ -25,6 +25,7 @@ const (
 	studioParentIDColumn  = "parent_id"
 	studioNameColumn      = "name"
 	studioImageBlobColumn = "image_blob"
+	studiosTagsTable      = "studios_tags"
 )
 
 type studioRow struct {
@@ -94,6 +95,7 @@ type studioRepositoryType struct {
 	repository
 
 	stashIDs stashIDRepository
+	tags     joinRepository
 
 	scenes    repository
 	images    repository
@@ -124,11 +126,21 @@ var (
 			tableName: galleryTable,
 			idColumn:  studioIDColumn,
 		},
+		tags: joinRepository{
+			repository: repository{
+				tableName: studiosTagsTable,
+				idColumn:  studioIDColumn,
+			},
+			fkColumn:     tagIDColumn,
+			foreignTable: tagTable,
+			orderBy:      "tags.name ASC",
+		},
 	}
 )
 
 type StudioStore struct {
 	blobJoinQueryBuilder
+	tagRelationshipStore
 
 	tableMgr *table
 }
@@ -139,6 +151,11 @@ func NewStudioStore(blobStore *BlobStore) *StudioStore {
 			blobStore: blobStore,
 			joinTable: studioTable,
 		},
+		tagRelationshipStore: tagRelationshipStore{
+			idRelationshipStore: idRelationshipStore{
+				joinTable: studiosTagsTableMgr,
+			},
+		},
 
 		tableMgr: studioTableMgr,
 	}
@@ -173,6 +190,10 @@ func (qb *StudioStore) Create(ctx context.Context, newObject *models.Studio) err
 		}
 	}
 
+	if err := qb.tagRelationshipStore.createRelationships(ctx, id, newObject.TagIDs); err != nil {
+		return err
+	}
+
 	if newObject.StashIDs.Loaded() {
 		if err := studiosStashIDsTableMgr.insertJoins(ctx, id, newObject.StashIDs.List()); err != nil {
 			return err
@@ -213,6 +234,10 @@ func (qb *StudioStore) UpdatePartial(ctx context.Context, input models.StudioPar
 		}
 	}
 
+	if err := qb.tagRelationshipStore.modifyRelationships(ctx, input.ID, input.TagIDs); err != nil {
+		return nil, err
+	}
+
 	if input.StashIDs != nil {
 		if err := studiosStashIDsTableMgr.modifyJoins(ctx, input.ID, input.StashIDs.StashIDs, input.StashIDs.Mode); err != nil {
 			return nil, err
@@ -237,6 +262,10 @@ func (qb *StudioStore) Update(ctx context.Context, updatedObject *models.Studio)
 		}
 	}
 
+	if err := qb.tagRelationshipStore.replaceRelationships(ctx, updatedObject.ID, updatedObject.TagIDs); err != nil {
+		return err
+	}
+
 	if updatedObject.StashIDs.Loaded() {
 		if err := studiosStashIDsTableMgr.replaceJoins(ctx, updatedObject.ID, updatedObject.StashIDs.List()); err != nil {
 			return err
@@ -538,6 +567,15 @@ func (qb *StudioStore) Query(ctx context.Context, studioFilter *models.StudioFil
 	return studios, countResult, nil
 }
 
+func (qb *StudioStore) QueryCount(ctx context.Context, studioFilter *models.StudioFilterType, findFilter *models.FindFilterType) (int, error) {
+	query, err := qb.makeQuery(ctx, studioFilter, findFilter)
+	if err != nil {
+		return 0, err
+	}
+
+	return query.executeCount(ctx)
+}
+
 var studioSortOptions = sortOptions{
 	"child_count",
 	"created_at",
@@ -569,6 +607,8 @@ func (qb *StudioStore) getStudioSort(findFilter *models.FindFilterType) (string,
 
 	sortQuery := ""
 	switch sort {
+	case "tag_count":
+		sortQuery += getCountSort(studioTable, studiosTagsTable, studioIDColumn, direction)
 	case "scenes_count":
 		sortQuery += getCountSort(studioTable, sceneTable, studioIDColumn, direction)
 	case "images_count":

pkg/sqlite/studio_filter.go

@@ -74,11 +74,13 @@ func (qb *studioFilterHandler) criterionHandler() criterionHandler {
 		},
 
 		qb.isMissingCriterionHandler(studioFilter.IsMissing),
+		qb.tagCountCriterionHandler(studioFilter.TagCount),
 		qb.sceneCountCriterionHandler(studioFilter.SceneCount),
 		qb.imageCountCriterionHandler(studioFilter.ImageCount),
 		qb.galleryCountCriterionHandler(studioFilter.GalleryCount),
 		qb.parentCriterionHandler(studioFilter.Parents),
 		qb.aliasCriterionHandler(studioFilter.Aliases),
+		qb.tagsCriterionHandler(studioFilter.Tags),
 		qb.childCountCriterionHandler(studioFilter.ChildCount),
 		&timestampCriterionHandler{studioFilter.CreatedAt, studioTable + ".created_at", nil},
 		&timestampCriterionHandler{studioFilter.UpdatedAt, studioTable + ".updated_at", nil},
@@ -161,6 +163,16 @@ func (qb *studioFilterHandler) galleryCountCriterionHandler(galleryCount *models
 	}
 }
 
+func (qb *studioFilterHandler) tagCountCriterionHandler(tagCount *models.IntCriterionInput) criterionHandlerFunc {
+	h := countCriterionHandlerBuilder{
+		primaryTable: studioTable,
+		joinTable:    studiosTagsTable,
+		primaryFK:    studioIDColumn,
+	}
+
+	return h.handler(tagCount)
+}
+
 func (qb *studioFilterHandler) parentCriterionHandler(parents *models.MultiCriterionInput) criterionHandlerFunc {
 	addJoinsFunc := func(f *filterBuilder) {
 		f.addLeftJoin("studios", "parent_studio", "parent_studio.id = studios.parent_id")
@@ -200,3 +212,18 @@ func (qb *studioFilterHandler) childCountCriterionHandler(childCount *models.Int
 		}
 	}
 }
+
+func (qb *studioFilterHandler) tagsCriterionHandler(tags *models.HierarchicalMultiCriterionInput) criterionHandlerFunc {
+	h := joinedHierarchicalMultiCriterionHandlerBuilder{
+		primaryTable: studioTable,
+		foreignTable: tagTable,
+		foreignFK:    "tag_id",
+
+		relationsTable: "tags_relations",
+		joinTable:      studiosTagsTable,
+		joinAs:         "studio_tag",
+		primaryFK:      studioIDColumn,
+	}
+
+	return h.handler(tags)
+}

pkg/sqlite/studio_test.go

@@ -704,6 +704,110 @@ func TestStudioQueryRating(t *testing.T) {
 	verifyStudiosRating(t, ratingCriterion)
 }
 
+func queryStudios(ctx context.Context, t *testing.T, studioFilter *models.StudioFilterType, findFilter *models.FindFilterType) []*models.Studio {
+	t.Helper()
+	studios, _, err := db.Studio.Query(ctx, studioFilter, findFilter)
+	if err != nil {
+		t.Errorf("Error querying studio: %s", err.Error())
+	}
+
+	return studios
+}
+
+func TestStudioQueryTags(t *testing.T) {
+	withTxn(func(ctx context.Context) error {
+		tagCriterion := models.HierarchicalMultiCriterionInput{
+			Value: []string{
+				strconv.Itoa(tagIDs[tagIdxWithStudio]),
+				strconv.Itoa(tagIDs[tagIdx1WithStudio]),
+			},
+			Modifier: models.CriterionModifierIncludes,
+		}
+
+		studioFilter := models.StudioFilterType{
+			Tags: &tagCriterion,
+		}
+
+		// ensure ids are correct
+		studios := queryStudios(ctx, t, &studioFilter, nil)
+		assert.Len(t, studios, 2)
+		for _, studio := range studios {
+			assert.True(t, studio.ID == studioIDs[studioIdxWithTag] || studio.ID == studioIDs[studioIdxWithTwoTags])
+		}
+
+		tagCriterion = models.HierarchicalMultiCriterionInput{
+			Value: []string{
+				strconv.Itoa(tagIDs[tagIdx1WithStudio]),
+				strconv.Itoa(tagIDs[tagIdx2WithStudio]),
+			},
+			Modifier: models.CriterionModifierIncludesAll,
+		}
+
+		studios = queryStudios(ctx, t, &studioFilter, nil)
+
+		assert.Len(t, studios, 1)
+		assert.Equal(t, sceneIDs[studioIdxWithTwoTags], studios[0].ID)
+
+		tagCriterion = models.HierarchicalMultiCriterionInput{
+			Value: []string{
+				strconv.Itoa(tagIDs[tagIdx1WithStudio]),
+			},
+			Modifier: models.CriterionModifierExcludes,
+		}
+
+		q := getSceneStringValue(studioIdxWithTwoTags, titleField)
+		findFilter := models.FindFilterType{
+			Q: &q,
+		}
+
+		studios = queryStudios(ctx, t, &studioFilter, &findFilter)
+		assert.Len(t, studios, 0)
+
+		return nil
+	})
+}
+
+func TestStudioQueryTagCount(t *testing.T) {
+	const tagCount = 1
+	tagCountCriterion := models.IntCriterionInput{
+		Value:    tagCount,
+		Modifier: models.CriterionModifierEquals,
+	}
+
+	verifyStudiosTagCount(t, tagCountCriterion)
+
+	tagCountCriterion.Modifier = models.CriterionModifierNotEquals
+	verifyStudiosTagCount(t, tagCountCriterion)
+
+	tagCountCriterion.Modifier = models.CriterionModifierGreaterThan
+	verifyStudiosTagCount(t, tagCountCriterion)
+
+	tagCountCriterion.Modifier = models.CriterionModifierLessThan
+	verifyStudiosTagCount(t, tagCountCriterion)
+}
+
+func verifyStudiosTagCount(t *testing.T, tagCountCriterion models.IntCriterionInput) {
+	withTxn(func(ctx context.Context) error {
+		sqb := db.Studio
+		studioFilter := models.StudioFilterType{
+			TagCount: &tagCountCriterion,
+		}
+
+		studios := queryStudios(ctx, t, &studioFilter, nil)
+		assert.Greater(t, len(studios), 0)
+
+		for _, studio := range studios {
+			ids, err := sqb.GetTagIDs(ctx, studio.ID)
+			if err != nil {
+				return err
+			}
+			verifyInt(t, len(ids), tagCountCriterion)
+		}
+
+		return nil
+	})
+}
+
 func verifyStudioQuery(t *testing.T, filter models.StudioFilterType, verifyFn func(ctx context.Context, s *models.Studio)) {
 	withTxn(func(ctx context.Context) error {
 		t.Helper()

pkg/sqlite/tables.go

@@ -34,6 +34,7 @@ var (
 	performersStashIDsJoinTable = goqu.T("performer_stash_ids")
 
 	studiosAliasesJoinTable  = goqu.T(studioAliasesTable)
+	studiosTagsJoinTable     = goqu.T(studiosTagsTable)
 	studiosStashIDsJoinTable = goqu.T("studio_stash_ids")
 
 	moviesURLsJoinTable = goqu.T(movieURLsTable)
@@ -294,6 +295,14 @@ var (
 		stringColumn: studiosAliasesJoinTable.Col(studioAliasColumn),
 	}
 
+	studiosTagsTableMgr = &joinTable{
+		table: table{
+			table:    studiosTagsJoinTable,
+			idColumn: studiosTagsJoinTable.Col(studioIDColumn),
+		},
+		fkColumn: studiosTagsJoinTable.Col(tagIDColumn),
+	}
+
 	studiosStashIDsTableMgr = &stashIDTable{
 		table: table{
 			table:    studiosStashIDsJoinTable,

pkg/sqlite/tag.go

@@ -448,6 +448,18 @@ func (qb *TagStore) FindBySceneMarkerID(ctx context.Context, sceneMarkerID int)
 	return qb.queryTags(ctx, query, args)
 }
 
+func (qb *TagStore) FindByStudioID(ctx context.Context, studioID int) ([]*models.Tag, error) {
+	query := `
+		SELECT tags.* FROM tags
+		LEFT JOIN studios_tags as studios_join on studios_join.tag_id = tags.id
+		WHERE studios_join.studio_id = ?
+		GROUP BY tags.id
+	`
+	query += qb.getDefaultTagSort()
+	args := []interface{}{studioID}
+	return qb.queryTags(ctx, query, args)
+}
+
 func (qb *TagStore) FindByName(ctx context.Context, name string, nocase bool) (*models.Tag, error) {
 	// query := "SELECT * FROM tags WHERE name = ?"
 	// if nocase {
@@ -628,6 +640,7 @@ var tagSortOptions = sortOptions{
 	"id",
 	"images_count",
 	"movies_count",
+	"studios_count",
 	"name",
 	"performers_count",
 	"random",
@@ -668,6 +681,8 @@ func (qb *TagStore) getTagSort(query *queryBuilder, findFilter *models.FindFilte
 		sortQuery += getCountSort(tagTable, galleriesTagsTable, tagIDColumn, direction)
 	case "performers_count":
 		sortQuery += getCountSort(tagTable, performersTagsTable, tagIDColumn, direction)
+	case "studios_count":
+		sortQuery += getCountSort(tagTable, studiosTagsTable, tagIDColumn, direction)
 	case "movies_count":
 		sortQuery += getCountSort(tagTable, moviesTagsTable, tagIDColumn, direction)
 	default:
@@ -767,6 +782,7 @@ func (qb *TagStore) Merge(ctx context.Context, source []int, destination int) er
 		galleriesTagsTable:   galleryIDColumn,
 		imagesTagsTable:      imageIDColumn,
 		"performers_tags":    "performer_id",
+		"studios_tags":       "studio_id",
 	}
 
 	args = append(args, destination)

pkg/sqlite/tag_filter.go

@@ -66,6 +66,7 @@ func (qb *tagFilterHandler) criterionHandler() criterionHandler {
 		qb.imageCountCriterionHandler(tagFilter.ImageCount),
 		qb.galleryCountCriterionHandler(tagFilter.GalleryCount),
 		qb.performerCountCriterionHandler(tagFilter.PerformerCount),
+		qb.studioCountCriterionHandler(tagFilter.StudioCount),
 		qb.movieCountCriterionHandler(tagFilter.MovieCount),
 		qb.markerCountCriterionHandler(tagFilter.MarkerCount),
 		qb.parentsCriterionHandler(tagFilter.Parents),
@@ -175,6 +176,17 @@ func (qb *tagFilterHandler) performerCountCriterionHandler(performerCount *model
 	}
 }
 
+func (qb *tagFilterHandler) studioCountCriterionHandler(studioCount *models.IntCriterionInput) criterionHandlerFunc {
+	return func(ctx context.Context, f *filterBuilder) {
+		if studioCount != nil {
+			f.addLeftJoin("studios_tags", "", "studios_tags.tag_id = tags.id")
+			clause, args := getIntCriterionWhereClause("count(distinct studios_tags.studio_id)", *studioCount)
+
+			f.addHaving(clause, args...)
+		}
+	}
+}
+
 func (qb *tagFilterHandler) movieCountCriterionHandler(movieCount *models.IntCriterionInput) criterionHandlerFunc {
 	return func(ctx context.Context, f *filterBuilder) {
 		if movieCount != nil {

pkg/sqlite/tag_test.go

@@ -230,6 +230,10 @@ func TestTagQuerySort(t *testing.T) {
 		tags = queryTags(ctx, t, sqb, nil, findFilter)
 		assert.Equal(tagIDs[tagIdx2WithPerformer], tags[0].ID)
 
+		sortBy = "studios_count"
+		tags = queryTags(ctx, t, sqb, nil, findFilter)
+		assert.Equal(tagIDs[tagIdx2WithStudio], tags[0].ID)
+
 		sortBy = "movies_count"
 		tags = queryTags(ctx, t, sqb, nil, findFilter)
 		assert.Equal(tagIDs[tagIdx1WithMovie], tags[0].ID)
@@ -569,6 +573,45 @@ func verifyTagPerformerCount(t *testing.T, imageCountCriterion models.IntCriteri
 	})
 }
 
+func TestTagQueryStudioCount(t *testing.T) {
+	countCriterion := models.IntCriterionInput{
+		Value:    1,
+		Modifier: models.CriterionModifierEquals,
+	}
+
+	verifyTagStudioCount(t, countCriterion)
+
+	countCriterion.Modifier = models.CriterionModifierNotEquals
+	verifyTagStudioCount(t, countCriterion)
+
+	countCriterion.Modifier = models.CriterionModifierLessThan
+	verifyTagStudioCount(t, countCriterion)
+
+	countCriterion.Value = 0
+	countCriterion.Modifier = models.CriterionModifierGreaterThan
+	verifyTagStudioCount(t, countCriterion)
+}
+
+func verifyTagStudioCount(t *testing.T, imageCountCriterion models.IntCriterionInput) {
+	withTxn(func(ctx context.Context) error {
+		qb := db.Tag
+		tagFilter := models.TagFilterType{
+			StudioCount: &imageCountCriterion,
+		}
+
+		tags, _, err := qb.Query(ctx, &tagFilter, nil)
+		if err != nil {
+			t.Errorf("Error querying tag: %s", err.Error())
+		}
+
+		for _, tag := range tags {
+			verifyInt(t, getTagStudioCount(tag.ID), imageCountCriterion)
+		}
+
+		return nil
+	})
+}
+
 func TestTagQueryParentCount(t *testing.T) {
 	countCriterion := models.IntCriterionInput{
 		Value:    1,
@@ -882,6 +925,9 @@ func TestTagMerge(t *testing.T) {
 			tagIdxWithPerformer,
 			tagIdx1WithPerformer,
 			tagIdx2WithPerformer,
+			tagIdxWithStudio,
+			tagIdx1WithStudio,
+			tagIdx2WithStudio,
 			tagIdxWithGallery,
 			tagIdx1WithGallery,
 			tagIdx2WithGallery,
@@ -970,6 +1016,14 @@ func TestTagMerge(t *testing.T) {
 
 		assert.Contains(performerTagIDs, destID)
 
+		// ensure studio points to new tag
+		studioTagIDs, err := db.Studio.GetTagIDs(ctx, studioIDs[studioIdxWithTwoTags])
+		if err != nil {
+			return err
+		}
+
+		assert.Contains(studioTagIDs, destID)
+
 		return nil
 	}); err != nil {
 		t.Error(err.Error())