rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add cookie banner and category dataloaders

Browse Source 
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>
This commit is contained in:
Émile Ré
2026-04-27 19:21:33 +04:00
parent d9b0d5f44d
commit b9835f0f25
7 changed files with 217 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
pkg/cookiebanner/service.go
View File

@@ -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,

48
pkg/coredata/cookie_banner.go
View File

@@ -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,

47
pkg/coredata/cookie_category.go
View File

@@ -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,

12
pkg/server/api/console/v1/cookie_banner_resolvers.go
View File

@@ -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))

8
pkg/server/api/console/v1/cookie_consent_record_resolvers.go
View File

@@ -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))

88
pkg/server/api/console/v1/dataloader/dataloader.go
View File

@@ -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]
Framework *dataloadgen.Loader[gid.GID, *coredata.Framework]
Control *dataloadgen.Loader[gid.GID, *coredata.Control]
Vendor *dataloadgen.Loader[gid.GID, *coredata.Vendor]
Document *dataloadgen.Loader[gid.GID, *coredata.Document]
Profile *dataloadgen.Loader[gid.GID, *coredata.MembershipProfile]
Risk *dataloadgen.Loader[gid.GID, *coredata.Risk]
Measure *dataloadgen.Loader[gid.GID, *coredata.Measure]
Task *dataloadgen.Loader[gid.GID, *coredata.Task]
File *dataloadgen.Loader[gid.GID, *coredata.File]
Report *dataloadgen.Loader[gid.GID, *coredata.Report]
Organization *dataloadgen.Loader[gid.GID, *coredata.Organization]
Framework *dataloadgen.Loader[gid.GID, *coredata.Framework]
Control *dataloadgen.Loader[gid.GID, *coredata.Control]
Vendor *dataloadgen.Loader[gid.GID, *coredata.Vendor]
Document *dataloadgen.Loader[gid.GID, *coredata.Document]
Profile *dataloadgen.Loader[gid.GID, *coredata.MembershipProfile]
Risk *dataloadgen.Loader[gid.GID, *coredata.Risk]
Measure *dataloadgen.Loader[gid.GID, *coredata.Measure]
Task *dataloadgen.Loader[gid.GID, *coredata.Task]
File *dataloadgen.Loader[gid.GID, *coredata.File]
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
iam *iam.Service
probo *probo.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),
Framework: dataloadgen.NewMappedLoader(f.fetchFrameworks),
Control: dataloadgen.NewMappedLoader(f.fetchControls),
Vendor: dataloadgen.NewMappedLoader(f.fetchVendors),
Document: dataloadgen.NewMappedLoader(f.fetchDocuments),
Profile: dataloadgen.NewMappedLoader(f.fetchProfiles),
Risk: dataloadgen.NewMappedLoader(f.fetchRisks),
Measure: dataloadgen.NewMappedLoader(f.fetchMeasures),
Task: dataloadgen.NewMappedLoader(f.fetchTasks),
File: dataloadgen.NewMappedLoader(f.fetchFiles),
Report: dataloadgen.NewMappedLoader(f.fetchReports),
Organization: dataloadgen.NewMappedLoader(f.fetchOrganizations),
Framework: dataloadgen.NewMappedLoader(f.fetchFrameworks),
Control: dataloadgen.NewMappedLoader(f.fetchControls),
Vendor: dataloadgen.NewMappedLoader(f.fetchVendors),
Document: dataloadgen.NewMappedLoader(f.fetchDocuments),
Profile: dataloadgen.NewMappedLoader(f.fetchProfiles),
Risk: dataloadgen.NewMappedLoader(f.fetchRisks),
Measure: dataloadgen.NewMappedLoader(f.fetchMeasures),
Task: dataloadgen.NewMappedLoader(f.fetchTasks),
File: dataloadgen.NewMappedLoader(f.fetchFiles),
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
}

2
pkg/server/api/console/v1/resolver.go
View File

@@ -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)