Add cookie banner and category dataloaders

Batch-load CookieBanner and CookieCategory entities via
dataloadgen instead of making individual service calls in
GraphQL resolvers, matching the existing dataloader pattern
used for organizations, frameworks, etc.

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

pkg/cookiebanner/service.go

@@ -692,6 +692,30 @@ func (s *Service) GetCookieBanner(
 	return &banner, nil
 }
 
+func (s *Service) GetCookieBannersByIDs(
+	ctx context.Context,
+	scope coredata.Scoper,
+	bannerIDs ...gid.GID,
+) (coredata.CookieBanners, error) {
+	var banners coredata.CookieBanners
+
+	err := s.pg.WithConn(
+		ctx,
+		func(ctx context.Context, conn pg.Querier) error {
+			if err := banners.LoadByIDs(ctx, conn, scope, bannerIDs); err != nil {
+				return fmt.Errorf("cannot load cookie banners by ids: %w", err)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return banners, nil
+}
+
 func (s *Service) GetActiveCookieBanner(
 	ctx context.Context,
 	bannerID gid.GID,
@@ -1072,6 +1096,30 @@ func (s *Service) GetCookieCategory(
 	return &category, nil
 }
 
+func (s *Service) GetCookieCategoriesByIDs(
+	ctx context.Context,
+	scope coredata.Scoper,
+	categoryIDs ...gid.GID,
+) (coredata.CookieCategories, error) {
+	var categories coredata.CookieCategories
+
+	err := s.pg.WithConn(
+		ctx,
+		func(ctx context.Context, conn pg.Querier) error {
+			if err := categories.LoadByIDs(ctx, conn, scope, categoryIDs); err != nil {
+				return fmt.Errorf("cannot load cookie categories by ids: %w", err)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return categories, nil
+}
+
 func (s *Service) ListCookieCategoriesForBanner(
 	ctx context.Context,
 	scope coredata.Scoper,

pkg/coredata/cookie_banner.go

@@ -234,6 +234,54 @@ LIMIT 1;
 	return nil
 }
 
+func (b *CookieBanners) LoadByIDs(
+	ctx context.Context,
+	conn pg.Querier,
+	scope Scoper,
+	bannerIDs []gid.GID,
+) error {
+	q := `
+SELECT
+	id,
+	organization_id,
+	name,
+	origin,
+	state,
+	privacy_policy_url,
+	cookie_policy_url,
+	consent_expiry_days,
+	consent_mode,
+	show_branding,
+	default_language,
+	created_at,
+	updated_at
+FROM
+	cookie_banners
+WHERE
+	%s
+	AND id = ANY(@banner_ids)
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment())
+
+	args := pgx.StrictNamedArgs{"banner_ids": bannerIDs}
+	maps.Copy(args, scope.SQLArguments())
+
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query cookie banners: %w", err)
+	}
+
+	banners, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CookieBanner])
+	if err != nil {
+		return fmt.Errorf("cannot collect cookie banners: %w", err)
+	}
+
+	*b = banners
+
+	return nil
+}
+
 func (b *CookieBanners) LoadByOrganizationID(
 	ctx context.Context,
 	conn pg.Querier,

pkg/coredata/cookie_category.go

@@ -146,6 +146,53 @@ LIMIT 1;
 	return nil
 }
 
+func (c *CookieCategories) LoadByIDs(
+	ctx context.Context,
+	conn pg.Querier,
+	scope Scoper,
+	categoryIDs []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 id = ANY(@category_ids)
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment())
+
+	args := pgx.StrictNamedArgs{"category_ids": categoryIDs}
+	maps.Copy(args, scope.SQLArguments())
+
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query cookie categories: %w", err)
+	}
+
+	categories, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[CookieCategory])
+	if err != nil {
+		return fmt.Errorf("cannot collect cookie categories: %w", err)
+	}
+
+	*c = categories
+
+	return nil
+}
+
 func (c *CookieCategories) LoadByCookieBannerID(
 	ctx context.Context,
 	conn pg.Querier,

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

@@ -29,11 +29,11 @@ func (r *cookieResolver) CookieCategory(ctx context.Context, obj *types.Cookie)
 		return nil, err
 	}
 
-	scope := coredata.NewScopeFromObjectID(obj.CookieCategory.ID)
+	loaders := dataloader.FromContext(ctx)
 
-	category, err := r.cookieBanner.GetCookieCategory(ctx, scope, obj.CookieCategory.ID)
+	category, err := loaders.CookieCategory.Load(ctx, obj.CookieCategory.ID)
 	if err != nil {
-		if errors.Is(err, cookiebanner.ErrCategoryNotFound) {
+		if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
 			return nil, nil
 		}
 		r.logger.ErrorCtx(ctx, "cannot get cookie category", log.Error(err))
@@ -277,11 +277,11 @@ func (r *cookieCategoryResolver) CookieBanner(ctx context.Context, obj *types.Co
 		return nil, err
 	}
 
-	scope := coredata.NewScopeFromObjectID(obj.CookieBanner.ID)
+	loaders := dataloader.FromContext(ctx)
 
-	banner, err := r.cookieBanner.GetCookieBanner(ctx, scope, obj.CookieBanner.ID)
+	banner, err := loaders.CookieBanner.Load(ctx, obj.CookieBanner.ID)
 	if err != nil {
-		if errors.Is(err, cookiebanner.ErrBannerNotFound) {
+		if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
 			return nil, nil
 		}
 		r.logger.ErrorCtx(ctx, "cannot get cookie banner", log.Error(err))

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

@@ -9,10 +9,12 @@ import (
 	"context"
 	"errors"
 
+	"github.com/vikstrous/dataloadgen"
 	"go.gearno.de/kit/log"
 	"go.probo.inc/probo/pkg/cookiebanner"
 	"go.probo.inc/probo/pkg/coredata"
 	"go.probo.inc/probo/pkg/probo"
+	"go.probo.inc/probo/pkg/server/api/console/v1/dataloader"
 	"go.probo.inc/probo/pkg/server/api/console/v1/schema"
 	"go.probo.inc/probo/pkg/server/api/console/v1/types"
 	"go.probo.inc/probo/pkg/server/gqlutils"
@@ -24,11 +26,11 @@ func (r *cookieConsentRecordResolver) CookieBanner(ctx context.Context, obj *typ
 		return nil, err
 	}
 
-	scope := coredata.NewScopeFromObjectID(obj.CookieBanner.ID)
+	loaders := dataloader.FromContext(ctx)
 
-	banner, err := r.cookieBanner.GetCookieBanner(ctx, scope, obj.CookieBanner.ID)
+	banner, err := loaders.CookieBanner.Load(ctx, obj.CookieBanner.ID)
 	if err != nil {
-		if errors.Is(err, cookiebanner.ErrBannerNotFound) {
+		if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
 			return nil, nil
 		}
 		r.logger.ErrorCtx(ctx, "cannot get cookie banner", log.Error(err))

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

@@ -20,6 +20,7 @@ import (
 	"net/http"
 
 	"github.com/vikstrous/dataloadgen"
+	"go.probo.inc/probo/pkg/cookiebanner"
 	"go.probo.inc/probo/pkg/coredata"
 	"go.probo.inc/probo/pkg/gid"
 	"go.probo.inc/probo/pkg/iam"
@@ -30,22 +31,25 @@ type (
 	ctxKey struct{ name string }
 
 	Loaders struct {
-		Organization *dataloadgen.Loader[gid.GID, *coredata.Organization]
+		Organization   *dataloadgen.Loader[gid.GID, *coredata.Organization]
-		Framework    *dataloadgen.Loader[gid.GID, *coredata.Framework]
+		Framework      *dataloadgen.Loader[gid.GID, *coredata.Framework]
-		Control      *dataloadgen.Loader[gid.GID, *coredata.Control]
+		Control        *dataloadgen.Loader[gid.GID, *coredata.Control]
-		Vendor       *dataloadgen.Loader[gid.GID, *coredata.Vendor]
+		Vendor         *dataloadgen.Loader[gid.GID, *coredata.Vendor]
-		Document     *dataloadgen.Loader[gid.GID, *coredata.Document]
+		Document       *dataloadgen.Loader[gid.GID, *coredata.Document]
-		Profile      *dataloadgen.Loader[gid.GID, *coredata.MembershipProfile]
+		Profile        *dataloadgen.Loader[gid.GID, *coredata.MembershipProfile]
-		Risk         *dataloadgen.Loader[gid.GID, *coredata.Risk]
+		Risk           *dataloadgen.Loader[gid.GID, *coredata.Risk]
-		Measure      *dataloadgen.Loader[gid.GID, *coredata.Measure]
+		Measure        *dataloadgen.Loader[gid.GID, *coredata.Measure]
-		Task         *dataloadgen.Loader[gid.GID, *coredata.Task]
+		Task           *dataloadgen.Loader[gid.GID, *coredata.Task]
-		File         *dataloadgen.Loader[gid.GID, *coredata.File]
+		File           *dataloadgen.Loader[gid.GID, *coredata.File]
-		Report       *dataloadgen.Loader[gid.GID, *coredata.Report]
+		Report         *dataloadgen.Loader[gid.GID, *coredata.Report]
+		CookieBanner   *dataloadgen.Loader[gid.GID, *coredata.CookieBanner]
+		CookieCategory *dataloadgen.Loader[gid.GID, *coredata.CookieCategory]
 	}
 
 	batchFetcher struct {
-		probo *probo.Service
+		probo        *probo.Service
-		iam   *iam.Service
+		iam          *iam.Service
+		cookieBanner *cookiebanner.Service
 	}
 )
 
@@ -55,11 +59,11 @@ func FromContext(ctx context.Context) *Loaders {
 	return ctx.Value(loadersKey).(*Loaders)
 }
 
-func NewMiddleware(proboSvc *probo.Service, iamSvc *iam.Service) func(http.Handler) http.Handler {
+func NewMiddleware(proboSvc *probo.Service, iamSvc *iam.Service, cookieBannerSvc *cookiebanner.Service) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(
 			func(w http.ResponseWriter, r *http.Request) {
-				f := &batchFetcher{probo: proboSvc, iam: iamSvc}
+				f := &batchFetcher{probo: proboSvc, iam: iamSvc, cookieBanner: cookieBannerSvc}
 				loaders := f.newLoaders()
 				ctx := context.WithValue(r.Context(), loadersKey, loaders)
 				next.ServeHTTP(w, r.WithContext(ctx))
@@ -70,17 +74,19 @@ func NewMiddleware(proboSvc *probo.Service, iamSvc *iam.Service) func(http.Handl
 
 func (f *batchFetcher) newLoaders() *Loaders {
 	return &Loaders{
-		Organization: dataloadgen.NewMappedLoader(f.fetchOrganizations),
+		Organization:   dataloadgen.NewMappedLoader(f.fetchOrganizations),
-		Framework:    dataloadgen.NewMappedLoader(f.fetchFrameworks),
+		Framework:      dataloadgen.NewMappedLoader(f.fetchFrameworks),
-		Control:      dataloadgen.NewMappedLoader(f.fetchControls),
+		Control:        dataloadgen.NewMappedLoader(f.fetchControls),
-		Vendor:       dataloadgen.NewMappedLoader(f.fetchVendors),
+		Vendor:         dataloadgen.NewMappedLoader(f.fetchVendors),
-		Document:     dataloadgen.NewMappedLoader(f.fetchDocuments),
+		Document:       dataloadgen.NewMappedLoader(f.fetchDocuments),
-		Profile:      dataloadgen.NewMappedLoader(f.fetchProfiles),
+		Profile:        dataloadgen.NewMappedLoader(f.fetchProfiles),
-		Risk:         dataloadgen.NewMappedLoader(f.fetchRisks),
+		Risk:           dataloadgen.NewMappedLoader(f.fetchRisks),
-		Measure:      dataloadgen.NewMappedLoader(f.fetchMeasures),
+		Measure:        dataloadgen.NewMappedLoader(f.fetchMeasures),
-		Task:         dataloadgen.NewMappedLoader(f.fetchTasks),
+		Task:           dataloadgen.NewMappedLoader(f.fetchTasks),
-		File:         dataloadgen.NewMappedLoader(f.fetchFiles),
+		File:           dataloadgen.NewMappedLoader(f.fetchFiles),
-		Report:       dataloadgen.NewMappedLoader(f.fetchReports),
+		Report:         dataloadgen.NewMappedLoader(f.fetchReports),
+		CookieBanner:   dataloadgen.NewMappedLoader(f.fetchCookieBanners),
+		CookieCategory: dataloadgen.NewMappedLoader(f.fetchCookieCategories),
 	}
 }
 
@@ -248,3 +254,33 @@ func (f *batchFetcher) fetchReports(ctx context.Context, keys []gid.GID) (map[gi
 	}
 	return result, nil
 }
+
+func (f *batchFetcher) fetchCookieBanners(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.CookieBanner, error) {
+	scope := coredata.NewScopeFromObjectID(keys[0])
+
+	banners, err := f.cookieBanner.GetCookieBannersByIDs(ctx, scope, keys...)
+	if err != nil {
+		return nil, fmt.Errorf("cannot batch load cookie banners: %w", err)
+	}
+
+	result := make(map[gid.GID]*coredata.CookieBanner, len(banners))
+	for _, v := range banners {
+		result[v.ID] = v
+	}
+	return result, nil
+}
+
+func (f *batchFetcher) fetchCookieCategories(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.CookieCategory, error) {
+	scope := coredata.NewScopeFromObjectID(keys[0])
+
+	categories, err := f.cookieBanner.GetCookieCategoriesByIDs(ctx, scope, keys...)
+	if err != nil {
+		return nil, fmt.Errorf("cannot batch load cookie categories: %w", err)
+	}
+
+	result := make(map[gid.GID]*coredata.CookieCategory, len(categories))
+	for _, v := range categories {
+		result[v.ID] = v
+	}
+	return result, nil
+}

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

@@ -93,7 +93,7 @@ func NewMux(
 		r.Use(authn.NewAPIKeyMiddleware(iamSvc, tokenSecret))
 		r.Use(authn.NewOAuth2AccessTokenMiddleware(iamSvc))
 		r.Use(authn.NewIdentityPresenceMiddleware())
-		r.Use(dataloader.NewMiddleware(proboSvc, iamSvc))
+		r.Use(dataloader.NewMiddleware(proboSvc, iamSvc, cookieBannerSvc))
 
 		r.Handle("/graphql", graphqlHandler)