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

Migrate pattern analysis worker to tracker patterns

Browse Source 
The worker now operates on TrackerPattern/DetectedTrackers
instead of CookiePattern/Cookies, with TrackerType included
in merge group keys to prevent cross-type merging.

Signed-off-by: Émile Ré <emile@getprobo.com>
This commit is contained in:
Émile Ré
2026-05-05 17:00:59 +04:00
parent 84120a07fc
commit b0610cac24
3 changed files with 139 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
pkg/cookiebanner/worker.go
View File

@@ -50,7 +50,7 @@ func NewPatternAnalysisWorker(
}
return worker.New(
"cookie-pattern-analysis-worker",
"tracker-pattern-analysis-worker",
h,
logger,
opts...,
@@ -94,13 +94,14 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
hasUncategorised = false
}
var exactPatterns coredata.CookiePatterns
var exactPatterns coredata.TrackerPatterns
if err := exactPatterns.LoadAllByCookieBannerID(
ctx,
tx,
scope,
banner.ID,
coredata.NewCookiePatternFilter(new(coredata.CookiePatternMatchTypeExact), nil, new(false)),
nil,
); err != nil {
return fmt.Errorf("cannot load exact patterns: %w", err)
}
@@ -113,11 +114,12 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
maxAge := mostCommonMaxAge(group)
source := bestSource(group)
prefixPattern := &coredata.CookiePattern{
ID: gid.New(banner.ID.TenantID(), coredata.CookiePatternEntityType),
prefixPattern := &coredata.TrackerPattern{
ID: gid.New(banner.ID.TenantID(), coredata.TrackerPatternEntityType),
OrganizationID: group[0].OrganizationID,
CookieBannerID: banner.ID,
CookieCategoryID: key.categoryID,
TrackerType: key.trackerType,
Pattern: key.prefix,
MatchType: coredata.CookiePatternMatchTypePrefix,
DisplayName: key.prefix + "*",
@@ -133,7 +135,7 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
return fmt.Errorf("cannot insert prefix pattern %q: %w", key.prefix, err)
}
if !inserted {
if err := prefixPattern.LoadByBannerIDAndPattern(ctx, tx, scope, banner.ID, key.prefix); err != nil {
if err := prefixPattern.LoadByBannerIDTypeAndPattern(ctx, tx, scope, banner.ID, key.trackerType, key.prefix); err != nil {
return fmt.Errorf("cannot load existing prefix pattern %q: %w", key.prefix, err)
}
@@ -143,9 +145,9 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
}
for _, exactPattern := range group {
var cookies coredata.Cookies
if err := cookies.RelinkByCookiePatternID(ctx, tx, scope, exactPattern.ID, prefixPattern.ID); err != nil {
return fmt.Errorf("cannot relink cookies from pattern %q: %w", exactPattern.Pattern, err)
var trackers coredata.DetectedTrackers
if err := trackers.RelinkByTrackerPatternID(ctx, tx, scope, exactPattern.ID, prefixPattern.ID); err != nil {
return fmt.Errorf("cannot relink detected trackers from pattern %q: %w", exactPattern.Pattern, err)
}
if err := exactPattern.Delete(ctx, tx, scope); err != nil {
@@ -170,7 +172,7 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
return fmt.Errorf("cannot adopt uncategorised patterns: %w", err)
}
var patterns coredata.CookiePatterns
var patterns coredata.TrackerPatterns
if err := patterns.RefreshLastMatchedAtByCookieBannerID(ctx, tx, scope, banner.ID); err != nil {
return fmt.Errorf("cannot refresh last_matched_at: %w", err)
}
@@ -187,25 +189,26 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
}
type mergeGroupKey struct {
categoryID gid.GID
prefix string
categoryID gid.GID
trackerType coredata.TrackerType
prefix string
}
func findMergeGroups(
patterns coredata.CookiePatterns,
patterns coredata.TrackerPatterns,
threshold int,
) map[mergeGroupKey][]*coredata.CookiePattern {
prefixCounts := make(map[mergeGroupKey][]*coredata.CookiePattern)
) map[mergeGroupKey][]*coredata.TrackerPattern {
prefixCounts := make(map[mergeGroupKey][]*coredata.TrackerPattern)
for _, p := range patterns {
for _, pfx := range separatorPrefixes(p.Pattern) {
key := mergeGroupKey{categoryID: p.CookieCategoryID, prefix: pfx}
key := mergeGroupKey{categoryID: p.CookieCategoryID, trackerType: p.TrackerType, prefix: pfx}
prefixCounts[key] = append(prefixCounts[key], p)
}
}
type candidate struct {
key mergeGroupKey
patterns []*coredata.CookiePattern
patterns []*coredata.TrackerPattern
}
var candidates []candidate
@@ -219,11 +222,11 @@ func findMergeGroups(
return len(candidates[i].key.prefix) > len(candidates[j].key.prefix)
})
assigned := make(map[*coredata.CookiePattern]bool)
groups := make(map[mergeGroupKey][]*coredata.CookiePattern)
assigned := make(map[*coredata.TrackerPattern]bool)
groups := make(map[mergeGroupKey][]*coredata.TrackerPattern)
for _, c := range candidates {
var unassigned []*coredata.CookiePattern
var unassigned []*coredata.TrackerPattern
for _, p := range c.patterns {
if !assigned[p] {
unassigned = append(unassigned, p)
@@ -253,16 +256,17 @@ func separatorPrefixes(name string) []string {
return prefixes
}
func bestSource(patterns []*coredata.CookiePattern) coredata.CookieSource {
func bestSource(patterns []*coredata.TrackerPattern) *coredata.CookieSource {
for _, p := range patterns {
if p.Source == coredata.CookieSourceScript {
return coredata.CookieSourceScript
if p.Source != nil && *p.Source == coredata.CookieSourceScript {
return p.Source
}
}
return coredata.CookieSourcePreExisting
src := coredata.CookieSourcePreExisting
return &src
}
func mostCommonMaxAge(patterns []*coredata.CookiePattern) *int {
func mostCommonMaxAge(patterns []*coredata.TrackerPattern) *int {
type key struct {
valid bool
val int
@@ -309,13 +313,14 @@ func (h *patternAnalysisHandler) adoptUncategorisedPatterns(
return false, fmt.Errorf("cannot load uncategorised category: %w", err)
}
var prefixPatterns coredata.CookiePatterns
var prefixPatterns coredata.TrackerPatterns
if err := prefixPatterns.LoadAllByCookieBannerID(
ctx,
tx,
scope,
banner.ID,
coredata.NewCookiePatternFilter(new(coredata.CookiePatternMatchTypePrefix), nil, new(false)),
nil,
); err != nil {
return false, fmt.Errorf("cannot load prefix patterns: %w", err)
}
@@ -329,22 +334,23 @@ func (h *patternAnalysisHandler) adoptUncategorisedPatterns(
})
exactMatchType := coredata.CookiePatternMatchTypeExact
var uncategorisedExact coredata.CookiePatterns
var uncategorisedExact coredata.TrackerPatterns
if err := uncategorisedExact.LoadAllByCookieBannerID(
ctx,
tx,
scope,
banner.ID,
coredata.NewCookiePatternFilter(&exactMatchType, &uncategorised.ID, new(false)),
nil,
); err != nil {
return false, fmt.Errorf("cannot load uncategorised exact patterns: %w", err)
}
adopted := false
for _, ep := range uncategorisedExact {
var match *coredata.CookiePattern
var match *coredata.TrackerPattern
for _, pp := range prefixPatterns {
if strings.HasPrefix(ep.Pattern, pp.Pattern) {
if ep.TrackerType == pp.TrackerType && strings.HasPrefix(ep.Pattern, pp.Pattern) {
match = pp
break
}
@@ -354,9 +360,9 @@ func (h *patternAnalysisHandler) adoptUncategorisedPatterns(
continue
}
var cookies coredata.Cookies
if err := cookies.RelinkByCookiePatternID(ctx, tx, scope, ep.ID, match.ID); err != nil {
return false, fmt.Errorf("cannot relink cookies from pattern %q: %w", ep.Pattern, err)
var trackers coredata.DetectedTrackers
if err := trackers.RelinkByTrackerPatternID(ctx, tx, scope, ep.ID, match.ID); err != nil {
return false, fmt.Errorf("cannot relink detected trackers from pattern %q: %w", ep.Pattern, err)
}
if err := ep.Delete(ctx, tx, scope); err != nil {

97
pkg/coredata/tracker_pattern.go
View File

@@ -118,6 +118,68 @@ LIMIT 1;
return nil
}
func (tp *TrackerPattern) LoadByBannerIDTypeAndPattern(
ctx context.Context,
conn pg.Querier,
scope Scoper,
cookieBannerID gid.GID,
trackerType TrackerType,
pattern string,
) error {
q := `
SELECT
id,
organization_id,
cookie_banner_id,
cookie_category_id,
tracker_type,
pattern,
match_type,
display_name,
description,
excluded,
max_age_seconds,
source,
last_matched_at,
created_at,
updated_at
FROM
tracker_patterns
WHERE
%s
AND cookie_banner_id = @cookie_banner_id
AND tracker_type = @tracker_type
AND pattern = @pattern
LIMIT 1;
`
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{
"cookie_banner_id": cookieBannerID,
"tracker_type": trackerType,
"pattern": pattern,
}
maps.Copy(args, scope.SQLArguments())
rows, err := conn.Query(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot query tracker patterns: %w", err)
}
p, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[TrackerPattern])
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return ErrResourceNotFound
}
return fmt.Errorf("cannot collect tracker pattern: %w", err)
}
*tp = p
return nil
}
func (tp *TrackerPattern) FindMatchingPattern(
ctx context.Context,
conn pg.Querier,
@@ -478,3 +540,38 @@ ORDER BY
return nil
}
func (tps *TrackerPatterns) RefreshLastMatchedAtByCookieBannerID(
ctx context.Context,
tx pg.Tx,
scope Scoper,
cookieBannerID gid.GID,
) error {
q := `
UPDATE tracker_patterns
SET
last_matched_at = sub.max_detected
FROM (
SELECT tracker_pattern_id, MAX(last_detected_at) AS max_detected
FROM detected_trackers
WHERE %[1]s AND cookie_banner_id = @cookie_banner_id
GROUP BY tracker_pattern_id
) sub
WHERE
tracker_patterns.id = sub.tracker_pattern_id
AND %[1]s
AND tracker_patterns.cookie_banner_id = @cookie_banner_id
`
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{"cookie_banner_id": cookieBannerID}
maps.Copy(args, scope.SQLArguments())
_, err := tx.Exec(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot refresh last_matched_at for banner tracker patterns: %w", err)
}
return nil
}

10
pkg/probod/probod.go
View File

@@ -656,12 +656,12 @@ func (impl *Implm) Run(
},
)
cookiePatternAnalysisWorker := cookiebanner.NewPatternAnalysisWorker(cookieBannerService, pgClient, l.Named("cookie-pattern-analysis-worker"))
cookiePatternAnalysisWorkerCtx, stopCookiePatternAnalysisWorker := context.WithCancel(context.Background())
trackerPatternAnalysisWorker := cookiebanner.NewPatternAnalysisWorker(cookieBannerService, pgClient, l.Named("tracker-pattern-analysis-worker"))
trackerPatternAnalysisWorkerCtx, stopTrackerPatternAnalysisWorker := context.WithCancel(context.Background())
wg.Go(
func() {
if err := cookiePatternAnalysisWorker.Run(cookiePatternAnalysisWorkerCtx); err != nil {
cancel(fmt.Errorf("cookie pattern analysis worker crashed: %w", err))
if err := trackerPatternAnalysisWorker.Run(trackerPatternAnalysisWorkerCtx); err != nil {
cancel(fmt.Errorf("tracker pattern analysis worker crashed: %w", err))
}
},
)
@@ -730,7 +730,7 @@ func (impl *Implm) Run(
stopTrustCenterServer()
stopWebhookSender()
stopESignService()
stopCookiePatternAnalysisWorker()
stopTrackerPatternAnalysisWorker()
stopMailingListWorker()
stopEvidenceDescriptionWorker()
stopDocumentPDFWorker()