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

Move common third-party resolver to thirdparty pkg

Browse Source 
resolveOrCreateCommonThirdParty lived as a package-level helper in the
tracker mapping worker, but the common pattern enrichment worker now
reuses it. Homing shared catalog logic in a mapping-named file made the
enrichment worker quietly depend on the mapping worker's file, and it is
not a mapping concern.

Move it to pkg/thirdparty as exported ResolveOrCreateCommonThirdParty,
decoupled from cookiebanner's TrackerMappingAgentResult (it now takes a
name and category) to avoid an import cycle. It stays a transaction-
scoped free function so both workers compose it into their own tx for
atomicity rather than receiving a service that owns its own connection.

Relocate the catalog dedup DB test alongside it.

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-05-29 19:35:17 +02:00
parent 9b03d199da
commit 8d5571cf91
5 changed files with 357 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
pkg/cookiebanner/common_pattern_enrichment_worker.go
View File

@@ -28,6 +28,7 @@ import (
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/thirdparty"
)
const defaultEnrichmentStaleAfter = 10 * time.Minute
@@ -118,20 +119,20 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
}
// Map before enriching: an unlinked pattern is run through the
// mapping agent first so a resolved vendor both seeds the enrichment
// mapping agent first so a confident vendor both seeds the enrichment
// prompt and gets linked. Attribution stays the mapping pipeline's
// job; the enricher only reuses it.
var thirdPartyID *gid.GID
// job; the enricher only reuses it. An already-linked pattern skips
// this entirely.
var attribution *TrackerMappingAgentResult
if cp.CommonThirdPartyID == nil {
id, name, err := h.identifyThirdParty(ctx, cp)
attribution, err = h.identifyThirdParty(ctx, cp)
if err != nil {
return err
}
thirdPartyID = id
if name != "" {
thirdPartyName = name
if attribution != nil {
thirdPartyName = attribution.ThirdPartyName
}
}
@@ -143,6 +144,19 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
// Resolve or create the catalog vendor only for an unlinked
// pattern; the mapping pipeline owns creation, so we reuse its
// name+slug dedup and never create a duplicate or override an
// existing link.
var thirdPartyID *gid.GID
if attribution != nil && cp.CommonThirdPartyID == nil {
thirdPartyID, err = thirdparty.ResolveOrCreateCommonThirdParty(ctx, tx, h.logger, attribution.ThirdPartyName, attribution.Category, nil)
if err != nil {
return fmt.Errorf("cannot resolve or create common third party: %w", err)
}
}
// A blank description is recorded as a terminal-for-now state:
// the row is marked enriched so the stale-recovery loop never
// re-queues it, but a later third-party link (mapping worker)
@@ -177,45 +191,6 @@ func (h *commonPatternEnrichmentHandler) Process(ctx context.Context, cp coredat
)
}
// resolveThirdPartyID maps the agent's returned company name to an
// existing catalog third party, but only when the pattern has none yet:
// the enrichment worker links, it never overrides an attribution the
// mapping pipeline already resolved. A name that matches no catalog row
// resolves to nil, so the worker never invents a third party.
func (h *commonPatternEnrichmentHandler) resolveThirdPartyID(
ctx context.Context,
cp coredata.CommonTrackerPattern,
thirdPartyName string,
) (*gid.GID, error) {
if cp.CommonThirdPartyID != nil || thirdPartyName == "" {
return nil, nil
}
var id *gid.GID
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
var party coredata.CommonThirdParty
if err := party.LoadByName(ctx, conn, thirdPartyName); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return nil
}
return err
}
id = &party.ID
return nil
},
); err != nil {
return nil, fmt.Errorf("cannot resolve common third party for enrichment: %w", err)
}
return id, nil
}
func (h *commonPatternEnrichmentHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
@@ -286,18 +261,17 @@ func (h *commonPatternEnrichmentHandler) research(
}
// identifyThirdParty reuses the tracker-mapping agent to attribute a
// vendor to an unlinked catalog pattern. It returns the resolved
// existing third party id and its name (for the enrichment prompt) only
// when the agent is confident and the name matches a catalog row;
// otherwise it returns nils so enrichment proceeds without a vendor. A
// failed agent run is best-effort and non-fatal, mirroring the mapping
// worker's identifyWithAgent.
// vendor to an unlinked catalog pattern. It performs no DB writes: it
// returns the confident attribution (name, category, confidence) or nil
// when the agent is unsure, leaving the caller to resolve or create the
// catalog row. A failed agent run is best-effort and non-fatal,
// mirroring the mapping worker's identifyWithAgent.
func (h *commonPatternEnrichmentHandler) identifyThirdParty(
ctx context.Context,
cp coredata.CommonTrackerPattern,
) (*gid.GID, string, error) {
) (*TrackerMappingAgentResult, error) {
if h.mappingAgent == nil {
return nil, "", nil
return nil, nil
}
prompt := buildCommonPatternIdentificationPrompt(cp)
@@ -323,22 +297,15 @@ func (h *commonPatternEnrichmentHandler) identifyThirdParty(
log.String("pattern", cp.Pattern),
)
return nil, "", nil
return nil, nil
}
name := strings.TrimSpace(result.Output.ThirdPartyName)
if name == "" || result.Output.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold {
return nil, "", nil
out := result.Output
out.ThirdPartyName = strings.TrimSpace(out.ThirdPartyName)
if out.ThirdPartyName == "" || out.ThirdPartyConfidence < agentThirdPartyConfidenceThreshold {
return nil, nil
}
id, err := h.resolveThirdPartyID(ctx, cp, name)
if err != nil {
return nil, "", err
}
if id == nil {
return nil, "", nil
}
return id, name, nil
return &out, nil
}

126
pkg/cookiebanner/common_pattern_enrichment_worker_test.go
View File

@@ -1,126 +0,0 @@
// 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 cookiebanner
import (
"context"
"io"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
func newEnrichmentHandler(client *pg.Client) *commonPatternEnrichmentHandler {
return &commonPatternEnrichmentHandler{
pg: client,
logger: log.NewLogger(log.WithOutput(io.Discard)),
}
}
// seedEnrichmentThirdParty inserts a collision-free catalog third party
// for the resolver to match against.
func seedEnrichmentThirdParty(t *testing.T, ctx context.Context, client *pg.Client) coredata.CommonThirdParty {
t.Helper()
now := time.Now().UTC().Truncate(time.Microsecond)
id := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType)
suffix := id.String()
party := coredata.CommonThirdParty{
ID: id,
Name: "Hotjar " + suffix,
Slug: "hotjar-" + suffix,
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
return party.Insert(ctx, tx)
}))
t.Cleanup(func() {
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
_, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, id)
return err
})
})
return party
}
// TestResolveThirdPartyID pins the enrichment worker's third-party
// resolution: it links the agent's returned company to an existing
// catalog row by name, but only when the pattern has no third party yet,
// and it never invents one for a name absent from the catalog.
func TestResolveThirdPartyID(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
h := newEnrichmentHandler(client)
party := seedEnrichmentThirdParty(t, ctx, client)
existingID := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType)
t.Run("links existing catalog third party when unset", func(t *testing.T) {
cp := coredata.CommonTrackerPattern{}
got, err := h.resolveThirdPartyID(ctx, cp, party.Name)
require.NoError(t, err)
require.NotNil(t, got)
assert.Equal(t, party.ID, *got)
})
t.Run("matches catalog name case-insensitively", func(t *testing.T) {
cp := coredata.CommonTrackerPattern{}
got, err := h.resolveThirdPartyID(ctx, cp, "hOtJaR "+party.ID.String())
require.NoError(t, err)
require.NotNil(t, got)
assert.Equal(t, party.ID, *got)
})
t.Run("does not resolve when pattern already linked", func(t *testing.T) {
cp := coredata.CommonTrackerPattern{CommonThirdPartyID: &existingID}
got, err := h.resolveThirdPartyID(ctx, cp, party.Name)
require.NoError(t, err)
assert.Nil(t, got, "must not override an existing third-party link")
})
t.Run("returns nil for a name absent from the catalog", func(t *testing.T) {
cp := coredata.CommonTrackerPattern{}
got, err := h.resolveThirdPartyID(ctx, cp, "Nonexistent Vendor "+party.ID.String())
require.NoError(t, err)
assert.Nil(t, got, "must not invent a third party")
})
t.Run("returns nil for an empty name", func(t *testing.T) {
cp := coredata.CommonTrackerPattern{}
got, err := h.resolveThirdPartyID(ctx, cp, "")
require.NoError(t, err)
assert.Nil(t, got)
})
}

68
pkg/cookiebanner/tracker_mapping_worker.go
View File

@@ -27,7 +27,6 @@ import (
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/slug"
"go.probo.inc/probo/pkg/thirdparty"
"go.probo.inc/probo/pkg/uri"
)
@@ -661,10 +660,12 @@ func (h *trackerMappingHandler) persistAgentIdentification(
tp coredata.TrackerPattern,
ident agentIdentification,
) (*catalogMatch, error) {
commonThirdPartyID, err := h.resolveOrCreateCommonThirdParty(
commonThirdPartyID, err := thirdparty.ResolveOrCreateCommonThirdParty(
ctx,
tx,
ident.result,
h.logger,
ident.result.ThirdPartyName,
ident.result.Category,
ident.domains,
)
if err != nil {
@@ -702,67 +703,6 @@ func (h *trackerMappingHandler) persistAgentIdentification(
}, nil
}
func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty(
ctx context.Context,
tx pg.Tx,
identification TrackerMappingAgentResult,
domains []string,
) (*gid.GID, error) {
var party coredata.CommonThirdParty
if err := party.LoadByName(ctx, tx, identification.ThirdPartyName); err == nil {
return &party.ID, nil
}
partySlug := slug.Make(identification.ThirdPartyName)
if partySlug == "" {
return nil, nil
}
if err := party.LoadBySlug(ctx, tx, partySlug); err == nil {
return &party.ID, nil
}
category := identification.Category
now := time.Now()
party = coredata.CommonThirdParty{
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType),
Name: identification.ThirdPartyName,
Slug: partySlug,
Category: category,
Certifications: []string{},
CreatedAt: now,
UpdatedAt: now,
}
if err := party.Insert(ctx, tx); err != nil {
return nil, fmt.Errorf("cannot create common third party: %w", err)
}
for _, domain := range domains {
domainRecord := coredata.CommonThirdPartyDomain{
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyDomainEntityType),
CommonThirdPartyID: party.ID,
Domain: domain,
CreatedAt: now,
UpdatedAt: now,
}
if _, err := domainRecord.Upsert(ctx, tx); err != nil {
return nil, fmt.Errorf("cannot create common third party domain: %w", err)
}
}
h.logger.InfoCtx(
ctx,
"created common third party from agent identification",
log.String("name", identification.ThirdPartyName),
log.String("category", category.String()),
)
return &party.ID, nil
}
// matchBySiblingOrigin finds other tracker patterns on the same banner
// that share initiator domains with the current pattern. Sharing an
// origin across multiple detected patterns is a strong indicator of the

94
pkg/thirdparty/resolver.go vendored Normal file
View File

@@ -0,0 +1,94 @@
// 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 thirdparty
import (
"context"
"fmt"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/slug"
)
// ResolveOrCreateCommonThirdParty links a named vendor to the global
// catalog, creating a row when none matches. Dedup is deterministic:
// exact name, then slug, before insert. Callers run inside their own
// transaction and pass the logger explicitly, so it is shared by the
// tracker mapping worker (which supplies observed domains) and the
// common pattern enrichment worker (which has none).
func ResolveOrCreateCommonThirdParty(
ctx context.Context,
tx pg.Tx,
logger *log.Logger,
name string,
category coredata.ThirdPartyCategory,
domains []string,
) (*gid.GID, error) {
var party coredata.CommonThirdParty
if err := party.LoadByName(ctx, tx, name); err == nil {
return &party.ID, nil
}
partySlug := slug.Make(name)
if partySlug == "" {
return nil, nil
}
if err := party.LoadBySlug(ctx, tx, partySlug); err == nil {
return &party.ID, nil
}
now := time.Now()
party = coredata.CommonThirdParty{
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType),
Name: name,
Slug: partySlug,
Category: category,
Certifications: []string{},
CreatedAt: now,
UpdatedAt: now,
}
if err := party.Insert(ctx, tx); err != nil {
return nil, fmt.Errorf("cannot create common third party: %w", err)
}
for _, domain := range domains {
domainRecord := coredata.CommonThirdPartyDomain{
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyDomainEntityType),
CommonThirdPartyID: party.ID,
Domain: domain,
CreatedAt: now,
UpdatedAt: now,
}
if _, err := domainRecord.Upsert(ctx, tx); err != nil {
return nil, fmt.Errorf("cannot create common third party domain: %w", err)
}
}
logger.InfoCtx(
ctx,
"created common third party from agent identification",
log.String("name", name),
log.String("category", category.String()),
)
return &party.ID, nil
}

224
pkg/thirdparty/resolver_test.go vendored Normal file
View File

@@ -0,0 +1,224 @@
// 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 thirdparty
import (
"context"
"io"
"net"
"net/url"
"os"
"testing"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/slug"
)
const testPgDSNEnvVar = "PROBO_TEST_PG_URL"
func newTestPgClient(t *testing.T) *pg.Client {
t.Helper()
dsn := os.Getenv(testPgDSNEnvVar)
if dsn == "" {
t.Skipf("skipping: %s not set (requires a migrated test database)", testPgDSNEnvVar)
}
u, err := url.Parse(dsn)
require.NoError(t, err, "invalid %s value", testPgDSNEnvVar)
opts := []pg.Option{pg.WithRegisterer(prometheus.NewRegistry())}
if u.Host != "" {
host := u.Host
if u.Port() == "" {
host = net.JoinHostPort(u.Hostname(), "5432")
}
opts = append(opts, pg.WithAddr(host))
}
if u.User != nil {
opts = append(opts, pg.WithUser(u.User.Username()))
if password, ok := u.User.Password(); ok {
opts = append(opts, pg.WithPassword(password))
}
}
if len(u.Path) > 1 {
opts = append(opts, pg.WithDatabase(u.Path[1:]))
}
client, err := pg.NewClient(opts...)
require.NoError(t, err)
t.Cleanup(func() {
client.Close()
})
return client
}
func discardLogger() *log.Logger {
return log.NewLogger(log.WithOutput(io.Discard))
}
// seedCatalogThirdParty inserts a catalog third party with the given
// name and slug and registers its cleanup.
func seedCatalogThirdParty(
t *testing.T,
ctx context.Context,
client *pg.Client,
name string,
slugValue string,
) coredata.CommonThirdParty {
t.Helper()
now := time.Now().UTC().Truncate(time.Microsecond)
party := coredata.CommonThirdParty{
ID: gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType),
Name: name,
Slug: slugValue,
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
return party.Insert(ctx, tx)
}))
t.Cleanup(func() {
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
_, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, party.ID)
return err
})
})
return party
}
// TestResolveOrCreateCommonThirdParty pins the catalog dedup that the
// mapping and enrichment workers reuse to link a vendor: an exact name
// match and a slug match both return the existing row, and a name absent
// from the catalog creates a fresh row (name, slug, category) rather than
// duplicating one.
func TestResolveOrCreateCommonThirdParty(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
logger := discardLogger()
token := slug.Make(gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType).String())
t.Run("returns existing row on name match", func(t *testing.T) {
name := "Hotjar " + token
party := seedCatalogThirdParty(t, ctx, client, name, slug.Make(name))
var got *gid.GID
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
id, err := ResolveOrCreateCommonThirdParty(
ctx,
tx,
logger,
name,
coredata.ThirdPartyCategoryAnalytics,
nil,
)
got = id
return err
}))
require.NotNil(t, got)
assert.Equal(t, party.ID, *got)
})
t.Run("returns existing row on slug match", func(t *testing.T) {
name := "Matomo " + token
party := seedCatalogThirdParty(t, ctx, client, name, slug.Make(name))
// A differently-spelled name that normalizes to the same slug
// must resolve to the existing row, not create a duplicate.
variant := "Matomo " + token + "!!!"
require.NotEqual(t, name, variant)
require.Equal(t, slug.Make(name), slug.Make(variant))
var got *gid.GID
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
id, err := ResolveOrCreateCommonThirdParty(
ctx,
tx,
logger,
variant,
coredata.ThirdPartyCategoryAnalytics,
nil,
)
got = id
return err
}))
require.NotNil(t, got)
assert.Equal(t, party.ID, *got)
})
t.Run("creates a new row when absent", func(t *testing.T) {
name := "Freshvendor " + token
expectedSlug := slug.Make(name)
var got *gid.GID
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
id, err := ResolveOrCreateCommonThirdParty(
ctx,
tx,
logger,
name,
coredata.ThirdPartyCategoryMarketing,
nil,
)
got = id
return err
}))
require.NotNil(t, got)
t.Cleanup(func() {
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
_, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, *got)
return err
})
})
var created coredata.CommonThirdParty
require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
return created.LoadByID(ctx, conn, *got)
}))
assert.Equal(t, name, created.Name)
assert.Equal(t, expectedSlug, created.Slug)
assert.Equal(t, coredata.ThirdPartyCategoryMarketing, created.Category)
})
}