diff --git a/pkg/cookiebanner/service.go b/pkg/cookiebanner/service.go index f8b8d7839..db07278dc 100644 --- a/pkg/cookiebanner/service.go +++ b/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, diff --git a/pkg/coredata/cookie_banner.go b/pkg/coredata/cookie_banner.go index 512649160..b5bd34a40 100644 --- a/pkg/coredata/cookie_banner.go +++ b/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, diff --git a/pkg/coredata/cookie_category.go b/pkg/coredata/cookie_category.go index b495e1896..4939e1f25 100644 --- a/pkg/coredata/cookie_category.go +++ b/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, diff --git a/pkg/server/api/console/v1/cookie_banner_resolvers.go b/pkg/server/api/console/v1/cookie_banner_resolvers.go index 698f291fa..83b7b4c83 100644 --- a/pkg/server/api/console/v1/cookie_banner_resolvers.go +++ b/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)) diff --git a/pkg/server/api/console/v1/cookie_consent_record_resolvers.go b/pkg/server/api/console/v1/cookie_consent_record_resolvers.go index edd3b2475..f7f5e7ddc 100644 --- a/pkg/server/api/console/v1/cookie_consent_record_resolvers.go +++ b/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)) diff --git a/pkg/server/api/console/v1/dataloader/dataloader.go b/pkg/server/api/console/v1/dataloader/dataloader.go index 30acd13a4..670c2d586 100644 --- a/pkg/server/api/console/v1/dataloader/dataloader.go +++ b/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] - 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 +} diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go index 878b6e904..9aa562055 100644 --- a/pkg/server/api/console/v1/resolver.go +++ b/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)