Unify cookie category queries with coredata filter

Replace duplicated LoadConsentCategoriesByCookieBannerID,
CountConsentCategoriesByCookieBannerID, and
LoadAllConsentCategoriesByCookieBannerID with a single
CookieCategoryFilter in pkg/coredata. The filter uses the
standard CASE WHEN idiom to optionally exclude a kind,
eliminating branching in the service layer.

Signed-off-by: Émile Ré <emile@probo.com>

pkg/cookiebanner/service.go

@@ -549,8 +549,10 @@ func (s *Service) ensureDraftVersionForBanner(
 		return nil, fmt.Errorf("cannot load cookie banner: %w", err)
 	}
 
+	consentFilter := coredata.NewCookieCategoryFilter(new(coredata.CookieCategoryKindUncategorised))
+
 	var categories coredata.CookieCategories
-	if err := categories.LoadAllConsentCategoriesByCookieBannerID(ctx, tx, scope, bannerID); err != nil {
+	if err := categories.LoadAllByCookieBannerID(ctx, tx, scope, bannerID, consentFilter); err != nil {
 		return nil, fmt.Errorf("cannot load cookie categories: %w", err)
 	}
 
@@ -1192,18 +1194,19 @@ func (s *Service) GetCookieCategoriesByIDs(
 	return categories, nil
 }
 
-func (s *Service) ListCookieCategoriesForBanner(
+func (s *Service) ListCategoriesForBanner(
 	ctx context.Context,
 	scope coredata.Scoper,
 	bannerID gid.GID,
 	cursor *page.Cursor[coredata.CookieCategoryOrderField],
+	filter *coredata.CookieCategoryFilter,
 ) (coredata.CookieCategories, error) {
 	var categories coredata.CookieCategories
 
 	err := s.pg.WithConn(
 		ctx,
 		func(ctx context.Context, conn pg.Querier) error {
-			if err := categories.LoadConsentCategoriesByCookieBannerID(ctx, conn, scope, bannerID, cursor); err != nil {
+			if err := categories.LoadByCookieBannerID(ctx, conn, scope, bannerID, cursor, filter); err != nil {
 				return fmt.Errorf("cannot list cookie categories: %w", err)
 			}
 
@@ -1217,10 +1220,11 @@ func (s *Service) ListCookieCategoriesForBanner(
 	return categories, nil
 }
 
-func (s *Service) CountCookieCategoriesForBanner(
+func (s *Service) CountCategoriesForBanner(
 	ctx context.Context,
 	scope coredata.Scoper,
 	bannerID gid.GID,
+	filter *coredata.CookieCategoryFilter,
 ) (int, error) {
 	var count int
 
@@ -1232,7 +1236,7 @@ func (s *Service) CountCookieCategoriesForBanner(
 				err        error
 			)
 
-			count, err = categories.CountConsentCategoriesByCookieBannerID(ctx, conn, scope, bannerID)
+			count, err = categories.CountByCookieBannerID(ctx, conn, scope, bannerID, filter)
 			if err != nil {
 				return fmt.Errorf("cannot count cookie categories: %w", err)
 			}
@@ -1637,8 +1641,10 @@ func (s *Service) GetActiveBannerConfig(
 				return fmt.Errorf("cannot get version snapshot: %w", err)
 			}
 
+			consentFilter := coredata.NewCookieCategoryFilter(new(coredata.CookieCategoryKindUncategorised))
+
 			var categories coredata.CookieCategories
-			if err := categories.LoadAllConsentCategoriesByCookieBannerID(ctx, conn, scope, banner.ID); err != nil {
+			if err := categories.LoadAllByCookieBannerID(ctx, conn, scope, banner.ID, consentFilter); err != nil {
 				return fmt.Errorf("cannot load cookie categories: %w", err)
 			}
 
@@ -2621,7 +2627,7 @@ func (s *Service) MoveTrackerPatternToCategory(
 	return &result, nil
 }
 
-func (s *Service) ListUncategorisedTrackerPatterns(
+func (s *Service) ListTrackerPatternsForBanner(
 	ctx context.Context,
 	scope coredata.Scoper,
 	bannerID gid.GID,
@@ -2633,8 +2639,8 @@ func (s *Service) ListUncategorisedTrackerPatterns(
 	err := s.pg.WithConn(
 		ctx,
 		func(ctx context.Context, conn pg.Querier) error {
-			if err := patterns.LoadUncategorisedByCookieBannerID(ctx, conn, scope, bannerID, cursor, filter); err != nil {
+			if err := patterns.LoadByCookieBannerID(ctx, conn, scope, bannerID, cursor, filter); err != nil {
-				return fmt.Errorf("cannot list uncategorised tracker patterns: %w", err)
+				return fmt.Errorf("cannot list tracker patterns for banner: %w", err)
 			}
 
 			return nil
@@ -2647,7 +2653,7 @@ func (s *Service) ListUncategorisedTrackerPatterns(
 	return patterns, nil
 }
 
-func (s *Service) CountUncategorisedTrackerPatterns(
+func (s *Service) CountTrackerPatternsForBanner(
 	ctx context.Context,
 	scope coredata.Scoper,
 	bannerID gid.GID,
@@ -2663,9 +2669,9 @@ func (s *Service) CountUncategorisedTrackerPatterns(
 				err      error
 			)
 
-			count, err = patterns.CountUncategorisedByCookieBannerID(ctx, conn, scope, bannerID, filter)
+			count, err = patterns.CountByCookieBannerID(ctx, conn, scope, bannerID, filter)
 			if err != nil {
-				return fmt.Errorf("cannot count uncategorised tracker patterns: %w", err)
+				return fmt.Errorf("cannot count tracker patterns for banner: %w", err)
 			}
 
 			return nil

pkg/coredata/cookie_category.go

@@ -202,6 +202,7 @@ func (c *CookieCategories) LoadByCookieBannerID(
 	scope Scoper,
 	cookieBannerID gid.GID,
 	cursor *page.Cursor[CookieCategoryOrderField],
+	filter *CookieCategoryFilter,
 ) error {
 	q := `
 SELECT
@@ -223,12 +224,14 @@ WHERE
 	%s
 	AND cookie_banner_id = @cookie_banner_id
 	AND %s
+	AND %s
 `
 
-	q = fmt.Sprintf(q, scope.SQLFragment(), cursor.SQLFragment())
+	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment(), cursor.SQLFragment())
 
 	args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID}
 	maps.Copy(args, scope.SQLArguments())
+	maps.Copy(args, filter.SQLArguments())
 	maps.Copy(args, cursor.SQLArguments())
 
 	rows, err := conn.Query(ctx, q, args)
@@ -251,6 +254,7 @@ func (c *CookieCategories) CountByCookieBannerID(
 	conn pg.Querier,
 	scope Scoper,
 	cookieBannerID gid.GID,
+	filter *CookieCategoryFilter,
 ) (int, error) {
 	q := `
 SELECT
@@ -260,101 +264,14 @@ FROM
 WHERE
 	%s
 	AND cookie_banner_id = @cookie_banner_id
-`
-
-	q = fmt.Sprintf(q, scope.SQLFragment())
-
-	args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID}
-	maps.Copy(args, scope.SQLArguments())
-
-	row := conn.QueryRow(ctx, q, args)
-
-	var count int
-	if err := row.Scan(&count); err != nil {
-		return 0, fmt.Errorf("cannot scan count: %w", err)
-	}
-
-	return count, nil
-}
-
-func (c *CookieCategories) LoadConsentCategoriesByCookieBannerID(
-	ctx context.Context,
-	conn pg.Querier,
-	scope Scoper,
-	cookieBannerID gid.GID,
-	cursor *page.Cursor[CookieCategoryOrderField],
-) error {
-	q := `
-SELECT
-	id,
-	organization_id,
-	cookie_banner_id,
-	name,
-	slug,
-	description,
-	kind,
-	rank,
-	gcm_consent_types,
-	posthog_consent,
-	created_at,
-	updated_at
-FROM
-	cookie_categories
-WHERE
-	%s
-	AND cookie_banner_id = @cookie_banner_id
-	AND kind != @excluded_kind
 	AND %s
 `
 
-	q = fmt.Sprintf(q, scope.SQLFragment(), cursor.SQLFragment())
+	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
 
-	args := pgx.StrictNamedArgs{
+	args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID}
-		"cookie_banner_id": cookieBannerID,
-		"excluded_kind":    CookieCategoryKindUncategorised,
-	}
-	maps.Copy(args, scope.SQLArguments())
-	maps.Copy(args, cursor.SQLArguments())
-
-	rows, err := conn.Query(ctx, q, args)
-	if err != nil {
-		return fmt.Errorf("cannot query consent cookie categories: %w", err)
-	}
-
-	categories, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CookieCategory])
-	if err != nil {
-		return fmt.Errorf("cannot collect consent cookie categories: %w", err)
-	}
-
-	*c = categories
-
-	return nil
-}
-
-func (c *CookieCategories) CountConsentCategoriesByCookieBannerID(
-	ctx context.Context,
-	conn pg.Querier,
-	scope Scoper,
-	cookieBannerID gid.GID,
-) (int, error) {
-	q := `
-SELECT
-	COUNT(id)
-FROM
-	cookie_categories
-WHERE
-	%s
-	AND cookie_banner_id = @cookie_banner_id
-	AND kind != @excluded_kind
-`
-
-	q = fmt.Sprintf(q, scope.SQLFragment())
-
-	args := pgx.StrictNamedArgs{
-		"cookie_banner_id": cookieBannerID,
-		"excluded_kind":    CookieCategoryKindUncategorised,
-	}
 	maps.Copy(args, scope.SQLArguments())
+	maps.Copy(args, filter.SQLArguments())
 
 	row := conn.QueryRow(ctx, q, args)
 
@@ -371,6 +288,7 @@ func (c *CookieCategories) LoadAllByCookieBannerID(
 	conn pg.Querier,
 	scope Scoper,
 	cookieBannerID gid.GID,
+	filter *CookieCategoryFilter,
 ) error {
 	q := `
 SELECT
@@ -391,14 +309,16 @@ FROM
 WHERE
 	%s
 	AND cookie_banner_id = @cookie_banner_id
+	AND %s
 ORDER BY
 	rank ASC, id ASC;
 `
 
-	q = fmt.Sprintf(q, scope.SQLFragment())
+	q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
 
 	args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID}
 	maps.Copy(args, scope.SQLArguments())
+	maps.Copy(args, filter.SQLArguments())
 
 	rows, err := conn.Query(ctx, q, args)
 	if err != nil {
@@ -415,61 +335,6 @@ ORDER BY
 	return nil
 }
 
-// LoadAllConsentCategoriesByCookieBannerID loads all categories except
-// UNCATEGORISED, which is an admin-side inbox never shown to visitors.
-func (c *CookieCategories) LoadAllConsentCategoriesByCookieBannerID(
-	ctx context.Context,
-	conn pg.Querier,
-	scope Scoper,
-	cookieBannerID gid.GID,
-) error {
-	q := `
-SELECT
-	id,
-	organization_id,
-	cookie_banner_id,
-	name,
-	slug,
-	description,
-	kind,
-	rank,
-	gcm_consent_types,
-	posthog_consent,
-	created_at,
-	updated_at
-FROM
-	cookie_categories
-WHERE
-	%s
-	AND cookie_banner_id = @cookie_banner_id
-	AND kind != @excluded_kind
-ORDER BY
-	rank ASC, id ASC;
-`
-
-	q = fmt.Sprintf(q, scope.SQLFragment())
-
-	args := pgx.StrictNamedArgs{
-		"cookie_banner_id": cookieBannerID,
-		"excluded_kind":    CookieCategoryKindUncategorised,
-	}
-	maps.Copy(args, scope.SQLArguments())
-
-	rows, err := conn.Query(ctx, q, args)
-	if err != nil {
-		return fmt.Errorf("cannot query consent cookie categories: %w", err)
-	}
-
-	categories, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CookieCategory])
-	if err != nil {
-		return fmt.Errorf("cannot collect consent cookie categories: %w", err)
-	}
-
-	*c = categories
-
-	return nil
-}
-
 func (c *CookieCategory) Insert(
 	ctx context.Context,
 	tx pg.Tx,

pkg/coredata/cookie_category_filter.go

@@ -0,0 +1,63 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package coredata
+
+import (
+	"github.com/jackc/pgx/v5"
+)
+
+type CookieCategoryFilter struct {
+	excludeKind *CookieCategoryKind
+}
+
+func NewCookieCategoryFilter(excludeKind *CookieCategoryKind) *CookieCategoryFilter {
+	return &CookieCategoryFilter{excludeKind: excludeKind}
+}
+
+func (f *CookieCategoryFilter) SQLFragment() string {
+	if f == nil {
+		return "TRUE"
+	}
+
+	return `(
+	CASE
+		WHEN @has_exclude_kind_filter::boolean = false THEN TRUE
+		WHEN @has_exclude_kind_filter::boolean = true THEN
+			kind != @filter_exclude_kind::cookie_category_kind
+		ELSE TRUE
+	END
+)`
+}
+
+func (f *CookieCategoryFilter) SQLArguments() pgx.StrictNamedArgs {
+	if f == nil {
+		return pgx.StrictNamedArgs{
+			"has_exclude_kind_filter": false,
+			"filter_exclude_kind":     nil,
+		}
+	}
+
+	args := pgx.StrictNamedArgs{
+		"has_exclude_kind_filter": false,
+		"filter_exclude_kind":     nil,
+	}
+
+	if f.excludeKind != nil {
+		args["has_exclude_kind_filter"] = true
+		args["filter_exclude_kind"] = string(*f.excludeKind)
+	}
+
+	return args
+}

pkg/server/api/console/v1/cookie_banner_resolvers.go

@@ -45,8 +45,8 @@ func (r *cookieBannerResolver) Organization(ctx context.Context, obj *types.Cook
 	return types.NewOrganization(organization), nil
 }
 
-// ConsentCategories is the resolver for the consentCategories field.
+// Categories is the resolver for the categories field.
-func (r *cookieBannerResolver) ConsentCategories(ctx context.Context, obj *types.CookieBanner, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.CookieCategoryOrderBy) (*types.CookieCategoryConnection, error) {
+func (r *cookieBannerResolver) Categories(ctx context.Context, obj *types.CookieBanner, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.CookieCategoryOrderBy, filter *types.CookieCategoryFilter) (*types.CookieCategoryConnection, error) {
 	if err := r.authorize(ctx, obj.ID, probo.ActionCookieCategoryList); err != nil {
 		return nil, err
 	}
@@ -65,7 +65,12 @@ func (r *cookieBannerResolver) ConsentCategories(ctx context.Context, obj *types
 	cursor := types.NewCursor(first, after, last, before, pageOrderBy)
 	scope := coredata.NewScopeFromObjectID(obj.ID)
 
-	categories, err := r.cookieBanner.ListCookieCategoriesForBanner(ctx, scope, obj.ID, cursor)
+	var cdFilter *coredata.CookieCategoryFilter
+	if filter != nil && filter.ExcludeKind != nil {
+		cdFilter = coredata.NewCookieCategoryFilter(filter.ExcludeKind)
+	}
+
+	categories, err := r.cookieBanner.ListCategoriesForBanner(ctx, scope, obj.ID, cursor, cdFilter)
 	if err != nil {
 		r.logger.ErrorCtx(ctx, "cannot list cookie categories", log.Error(err))
 		return nil, gqlutils.Internal(ctx)
@@ -73,7 +78,7 @@ func (r *cookieBannerResolver) ConsentCategories(ctx context.Context, obj *types
 
 	p := page.NewPage(categories, cursor)
 
-	return types.NewCookieCategoryConnection(p, r, obj.ID), nil
+	return types.NewCookieCategoryConnectionWithFilter(p, r, obj.ID, cdFilter), nil
 }
 
 // Translations is the resolver for the translations field.
@@ -180,8 +185,8 @@ func (r *cookieBannerResolver) ConsentRecords(ctx context.Context, obj *types.Co
 	return types.NewCookieConsentRecordConnection(p, r, obj.ID, coredataFilter), nil
 }
 
-// UncategorisedTrackerPatterns is the resolver for the uncategorisedTrackerPatterns field.
+// TrackerPatterns is the resolver for the trackerPatterns field.
-func (r *cookieBannerResolver) UncategorisedTrackerPatterns(ctx context.Context, obj *types.CookieBanner, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.TrackerPatternOrderBy, filter *types.TrackerPatternFilter) (*types.TrackerPatternConnection, error) {
+func (r *cookieBannerResolver) TrackerPatterns(ctx context.Context, obj *types.CookieBanner, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.TrackerPatternOrderBy, filter *types.TrackerPatternFilter) (*types.TrackerPatternConnection, error) {
 	if err := r.authorize(ctx, obj.ID, probo.ActionTrackerPatternList); err != nil {
 		return nil, err
 	}
@@ -202,12 +207,13 @@ func (r *cookieBannerResolver) UncategorisedTrackerPatterns(ctx context.Context,
 
 	coredataFilter := coredata.NewTrackerPatternFilter(nil, nil, nil)
 	if filter != nil {
+		coredataFilter = coredata.NewTrackerPatternFilter(nil, filter.CookieCategoryID, nil)
 		coredataFilter = coredataFilter.WithQuery(filter.Query).WithSource(filter.Source).WithTrackerType(filter.TrackerType)
 	}
 
-	patterns, err := r.cookieBanner.ListUncategorisedTrackerPatterns(ctx, scope, obj.ID, cursor, coredataFilter)
+	patterns, err := r.cookieBanner.ListTrackerPatternsForBanner(ctx, scope, obj.ID, cursor, coredataFilter)
 	if err != nil {
-		r.logger.ErrorCtx(ctx, "cannot list uncategorised tracker patterns", log.Error(err))
+		r.logger.ErrorCtx(ctx, "cannot list tracker patterns", log.Error(err))
 		return nil, gqlutils.Internal(ctx)
 	}
 
@@ -418,7 +424,7 @@ func (r *cookieCategoryConnectionResolver) TotalCount(ctx context.Context, obj *
 
 	scope := coredata.NewScopeFromObjectID(obj.ParentID)
 
-	count, err := r.cookieBanner.CountCookieCategoriesForBanner(ctx, scope, obj.ParentID)
+	count, err := r.cookieBanner.CountCategoriesForBanner(ctx, scope, obj.ParentID, obj.Filter)
 	if err != nil {
 		r.logger.ErrorCtx(ctx, "cannot count cookie categories", log.Error(err))
 		return 0, gqlutils.Internal(ctx)
@@ -1291,10 +1297,11 @@ func (r *trackerPatternConnectionResolver) TotalCount(ctx context.Context, obj *
 	default:
 		filter := coredata.NewTrackerPatternFilter(nil, nil, nil)
 		if obj.Filter != nil {
+			filter = coredata.NewTrackerPatternFilter(nil, obj.Filter.CookieCategoryID, nil)
 			filter = filter.WithQuery(obj.Filter.Query).WithSource(obj.Filter.Source).WithTrackerType(obj.Filter.TrackerType)
 		}
 
-		count, err = r.cookieBanner.CountUncategorisedTrackerPatterns(ctx, scope, obj.ParentID, filter)
+		count, err = r.cookieBanner.CountTrackerPatternsForBanner(ctx, scope, obj.ParentID, filter)
 	}
 
 	if err != nil {

pkg/server/api/console/v1/types/cookie_category.go

@@ -23,6 +23,10 @@ import (
 type (
 	CookieCategoryOrderBy OrderBy[coredata.CookieCategoryOrderField]
 
+	CookieCategoryFilter struct {
+		ExcludeKind *coredata.CookieCategoryKind
+	}
+
 	CookieCategoryConnection struct {
 		TotalCount int
 		Edges      []*CookieCategoryEdge
@@ -30,6 +34,7 @@ type (
 
 		Resolver any
 		ParentID gid.GID
+		Filter   *coredata.CookieCategoryFilter
 	}
 )
 
@@ -53,6 +58,18 @@ func NewCookieCategoryConnection(
 	}
 }
 
+func NewCookieCategoryConnectionWithFilter(
+	p *page.Page[*coredata.CookieCategory, coredata.CookieCategoryOrderField],
+	parentType any,
+	parentID gid.GID,
+	filter *coredata.CookieCategoryFilter,
+) *CookieCategoryConnection {
+	conn := NewCookieCategoryConnection(p, parentType, parentID)
+	conn.Filter = filter
+
+	return conn
+}
+
 func NewCookieCategoryEdge(c *coredata.CookieCategory, orderBy coredata.CookieCategoryOrderField) *CookieCategoryEdge {
 	return &CookieCategoryEdge{
 		Cursor: c.CursorKey(orderBy),

pkg/server/api/mcp/v1/schema.resolvers.go

@@ -5053,7 +5053,7 @@ func (r *Resolver) ListCookieCategoriesTool(ctx context.Context, req *mcp.CallTo
 	scope := coredata.NewScopeFromObjectID(input.CookieBannerID)
 	cursor := types.NewCursor(input.Size, input.Cursor, page.OrderBy[coredata.CookieCategoryOrderField]{Field: coredata.CookieCategoryOrderFieldRank, Direction: page.OrderDirectionAsc})
 
-	categories, err := r.cookieBanner.ListCookieCategoriesForBanner(ctx, scope, input.CookieBannerID, cursor)
+	categories, err := r.cookieBanner.ListCategoriesForBanner(ctx, scope, input.CookieBannerID, cursor, coredata.NewCookieCategoryFilter(new(coredata.CookieCategoryKindUncategorised)))
 	if err != nil {
 		panic(fmt.Errorf("cannot list cookie categories: %w", err))
 	}