Add common third party reenrich and stats CLIs
Add `proboctl common-third-party reenrich` to re-arm the async enrichment worker for selected catalog rows, and `stats` to summarize the catalog by enrichment state and last run status. Rows are selected verbatim via --id/--slug or across the catalog via --category/--keyword/--state/--status, gated by --dry-run and --yes. Extend `list` with --state/--status filters and STATE/STATUS columns, and `show` with enrichment state, attempts, last run status, error, per-field provenance, and discovered domains. Back these with CommonThirdPartyFilter state/status/IDs filters plus CommonThirdParties.LoadAllIDs and RequestEnrichmentByIDs. The latter stamps enrichment_requested_at and resets the attempt counter while preserving the existing payload, so the worker merge keeps curated and human-edited provenance. Also simplify exactLabelMatch to use slices.Contains. Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
@@ -655,6 +655,42 @@ LIMIT 20
|
||||
return nil
|
||||
}
|
||||
|
||||
// LoadAllIDs returns the IDs of every common third party matching the
|
||||
// filter, ignoring pagination. It is the selection primitive behind bulk
|
||||
// operator actions such as re-arming enrichment across a filtered set.
|
||||
func (t *CommonThirdParties) LoadAllIDs(
|
||||
ctx context.Context,
|
||||
conn pg.Querier,
|
||||
filter *CommonThirdPartyFilter,
|
||||
) ([]gid.GID, error) {
|
||||
q := `
|
||||
SELECT
|
||||
id
|
||||
FROM
|
||||
common_third_parties
|
||||
WHERE
|
||||
%s
|
||||
ORDER BY name ASC
|
||||
`
|
||||
|
||||
q = fmt.Sprintf(q, filter.SQLFragment())
|
||||
|
||||
args := pgx.StrictNamedArgs{}
|
||||
maps.Copy(args, filter.SQLArguments())
|
||||
|
||||
rows, err := conn.Query(ctx, q, args)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot query common third party ids: %w", err)
|
||||
}
|
||||
|
||||
ids, err := pgx.CollectRows(rows, pgx.RowTo[gid.GID])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot collect common third party ids: %w", err)
|
||||
}
|
||||
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func (t CommonThirdParty) UpdateLogoFileID(
|
||||
ctx context.Context,
|
||||
conn pg.Tx,
|
||||
@@ -1003,3 +1039,40 @@ WHERE
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RequestEnrichmentByIDs re-arms enrichment on the given common third
|
||||
// parties by stamping enrichment_requested_at, which is the only column
|
||||
// the enrichment worker claims on. It resets enrichment_attempts to 0 so
|
||||
// the row gets a fresh retry budget: the claim path bumps the counter on
|
||||
// every run, and without a reset a row near the max-attempts ceiling
|
||||
// would not be re-armed by stale recovery if a re-run crashed. The
|
||||
// enrichment payload is left in place so the worker's merge keeps prior
|
||||
// per-field provenance (it only overwrites fields it owns, never curated
|
||||
// seed data or human edits). Already-enriched rows are re-processed too.
|
||||
// Returns the number of rows re-queued.
|
||||
func (t *CommonThirdParties) RequestEnrichmentByIDs(
|
||||
ctx context.Context,
|
||||
tx pg.Tx,
|
||||
ids []gid.GID,
|
||||
) (int64, error) {
|
||||
q := `
|
||||
UPDATE common_third_parties
|
||||
SET
|
||||
enrichment_requested_at = NOW(),
|
||||
enrichment_attempts = 0,
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
id = ANY(@ids)
|
||||
`
|
||||
|
||||
args := pgx.StrictNamedArgs{
|
||||
"ids": ids,
|
||||
}
|
||||
|
||||
result, err := tx.Exec(ctx, q, args)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("cannot request common third party enrichment: %w", err)
|
||||
}
|
||||
|
||||
return result.RowsAffected(), nil
|
||||
}
|
||||
|
||||
@@ -15,19 +15,83 @@
|
||||
package coredata
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"go.probo.inc/probo/pkg/gid"
|
||||
)
|
||||
|
||||
// CommonThirdPartyEnrichmentState is a synthetic filter over the
|
||||
// enrichment_requested_at / enrichment columns. It is not a stored
|
||||
// column; it classifies a row's position in the enrichment lifecycle.
|
||||
// Unlike common tracker patterns, a common third party has no
|
||||
// enriched_at column: a row is "enriched" once it carries an enrichment
|
||||
// payload (Process always writes one, even on a no-result run).
|
||||
type CommonThirdPartyEnrichmentState string
|
||||
|
||||
const (
|
||||
// CommonThirdPartyEnrichmentStateQueued: a row armed for the
|
||||
// enrichment worker (enrichment_requested_at IS NOT NULL).
|
||||
CommonThirdPartyEnrichmentStateQueued CommonThirdPartyEnrichmentState = "QUEUED"
|
||||
// CommonThirdPartyEnrichmentStateEnriched: a row whose enrichment has
|
||||
// completed (enrichment IS NOT NULL) and is not re-queued.
|
||||
CommonThirdPartyEnrichmentStateEnriched CommonThirdPartyEnrichmentState = "ENRICHED"
|
||||
// CommonThirdPartyEnrichmentStateUnenriched: a row never enriched and
|
||||
// not currently queued.
|
||||
CommonThirdPartyEnrichmentStateUnenriched CommonThirdPartyEnrichmentState = "UNENRICHED"
|
||||
)
|
||||
|
||||
func (s CommonThirdPartyEnrichmentState) IsValid() bool {
|
||||
switch s {
|
||||
case
|
||||
CommonThirdPartyEnrichmentStateQueued,
|
||||
CommonThirdPartyEnrichmentStateEnriched,
|
||||
CommonThirdPartyEnrichmentStateUnenriched:
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func (s CommonThirdPartyEnrichmentState) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
func (s CommonThirdPartyEnrichmentState) MarshalText() ([]byte, error) {
|
||||
return []byte(s.String()), nil
|
||||
}
|
||||
|
||||
func (s *CommonThirdPartyEnrichmentState) UnmarshalText(text []byte) error {
|
||||
val := CommonThirdPartyEnrichmentState(text)
|
||||
if !val.IsValid() {
|
||||
return fmt.Errorf("invalid CommonThirdPartyEnrichmentState value: %q", string(text))
|
||||
}
|
||||
|
||||
*s = val
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type CommonThirdPartyFilter struct {
|
||||
name *string
|
||||
category *ThirdPartyCategory
|
||||
keyword *string
|
||||
ids []gid.GID
|
||||
name *string
|
||||
category *ThirdPartyCategory
|
||||
keyword *string
|
||||
state *CommonThirdPartyEnrichmentState
|
||||
enrichmentStatus *string
|
||||
}
|
||||
|
||||
func NewCommonThirdPartyFilter(name *string) *CommonThirdPartyFilter {
|
||||
return &CommonThirdPartyFilter{name: name}
|
||||
}
|
||||
|
||||
// WithIDs restricts the result to the given common third party IDs. A
|
||||
// non-nil but empty slice matches nothing.
|
||||
func (f *CommonThirdPartyFilter) WithIDs(ids []gid.GID) *CommonThirdPartyFilter {
|
||||
f.ids = ids
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *CommonThirdPartyFilter) WithCategory(category *ThirdPartyCategory) *CommonThirdPartyFilter {
|
||||
f.category = category
|
||||
return f
|
||||
@@ -38,8 +102,27 @@ func (f *CommonThirdPartyFilter) WithKeyword(keyword *string) *CommonThirdPartyF
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *CommonThirdPartyFilter) WithState(state *CommonThirdPartyEnrichmentState) *CommonThirdPartyFilter {
|
||||
f.state = state
|
||||
return f
|
||||
}
|
||||
|
||||
// WithEnrichmentStatus filters on the run-level status recorded in the
|
||||
// enrichment payload (done, partial, failed). Rows with no payload never
|
||||
// match.
|
||||
func (f *CommonThirdPartyFilter) WithEnrichmentStatus(status *string) *CommonThirdPartyFilter {
|
||||
f.enrichmentStatus = status
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *CommonThirdPartyFilter) SQLFragment() string {
|
||||
return `(
|
||||
CASE
|
||||
WHEN @filter_ids::text[] IS NOT NULL THEN
|
||||
id = ANY(@filter_ids)
|
||||
ELSE TRUE
|
||||
END
|
||||
AND
|
||||
CASE
|
||||
WHEN @filter_name::text IS NOT NULL THEN
|
||||
name ILIKE '%' || @filter_name || '%'
|
||||
@@ -58,14 +141,38 @@ func (f *CommonThirdPartyFilter) SQLFragment() string {
|
||||
OR slug ILIKE '%' || @filter_keyword || '%')
|
||||
ELSE TRUE
|
||||
END
|
||||
AND
|
||||
CASE
|
||||
WHEN @filter_state_queued::boolean THEN enrichment_requested_at IS NOT NULL
|
||||
WHEN @filter_state_enriched::boolean THEN
|
||||
enrichment_requested_at IS NULL AND enrichment IS NOT NULL
|
||||
WHEN @filter_state_unenriched::boolean THEN
|
||||
enrichment_requested_at IS NULL AND enrichment IS NULL
|
||||
ELSE TRUE
|
||||
END
|
||||
AND
|
||||
CASE
|
||||
WHEN @filter_enrichment_status::text IS NOT NULL THEN
|
||||
enrichment->>'status' = @filter_enrichment_status
|
||||
ELSE TRUE
|
||||
END
|
||||
)`
|
||||
}
|
||||
|
||||
func (f *CommonThirdPartyFilter) SQLArguments() pgx.StrictNamedArgs {
|
||||
args := pgx.StrictNamedArgs{
|
||||
"filter_name": nil,
|
||||
"filter_category": nil,
|
||||
"filter_keyword": nil,
|
||||
"filter_ids": nil,
|
||||
"filter_name": nil,
|
||||
"filter_category": nil,
|
||||
"filter_keyword": nil,
|
||||
"filter_state_queued": false,
|
||||
"filter_state_enriched": false,
|
||||
"filter_state_unenriched": false,
|
||||
"filter_enrichment_status": nil,
|
||||
}
|
||||
|
||||
if f.ids != nil {
|
||||
args["filter_ids"] = f.ids
|
||||
}
|
||||
|
||||
if f.name != nil {
|
||||
@@ -80,5 +187,20 @@ func (f *CommonThirdPartyFilter) SQLArguments() pgx.StrictNamedArgs {
|
||||
args["filter_keyword"] = *f.keyword
|
||||
}
|
||||
|
||||
if f.state != nil {
|
||||
switch *f.state {
|
||||
case CommonThirdPartyEnrichmentStateQueued:
|
||||
args["filter_state_queued"] = true
|
||||
case CommonThirdPartyEnrichmentStateEnriched:
|
||||
args["filter_state_enriched"] = true
|
||||
case CommonThirdPartyEnrichmentStateUnenriched:
|
||||
args["filter_state_unenriched"] = true
|
||||
}
|
||||
}
|
||||
|
||||
if f.enrichmentStatus != nil {
|
||||
args["filter_enrichment_status"] = *f.enrichmentStatus
|
||||
}
|
||||
|
||||
return args
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ package commonthirdparty
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
@@ -27,23 +28,100 @@ import (
|
||||
"go.probo.inc/probo/pkg/proboctl/cmdutil"
|
||||
)
|
||||
|
||||
// NewCmdCommonThirdParty is the entry point for inspecting the global
|
||||
// common third party catalog.
|
||||
// NewCmdCommonThirdParty is the entry point for inspecting and
|
||||
// re-enriching the global common third party catalog.
|
||||
func NewCmdCommonThirdParty(f *cmdutil.Factory) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "common-third-party <command>",
|
||||
Aliases: []string{"ctp3"},
|
||||
Short: "Inspect the global common third party catalog",
|
||||
Short: "Inspect and re-enrich the global common third party catalog",
|
||||
}
|
||||
|
||||
cmd.AddCommand(newCmdList(f))
|
||||
cmd.AddCommand(newCmdShow(f))
|
||||
cmd.AddCommand(newCmdDomains(f))
|
||||
cmd.AddCommand(newCmdUpsert(f))
|
||||
cmd.AddCommand(newCmdReenrich(f))
|
||||
cmd.AddCommand(newCmdStats(f))
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
// enrichmentState classifies a common third party's position in the
|
||||
// enrichment lifecycle for display. A row is "enriched" once it carries
|
||||
// an enrichment payload; there is no enriched_at column.
|
||||
func enrichmentState(p *coredata.CommonThirdParty) string {
|
||||
switch {
|
||||
case p.EnrichmentRequestedAt != nil:
|
||||
return "queued"
|
||||
case len(p.Enrichment) > 0:
|
||||
return "enriched"
|
||||
default:
|
||||
return "unenriched"
|
||||
}
|
||||
}
|
||||
|
||||
// enrichmentStatus returns the run-level status recorded in the
|
||||
// enrichment payload (done, partial, failed), or an empty string when
|
||||
// the row has never been enriched or the payload is malformed.
|
||||
func enrichmentStatus(p *coredata.CommonThirdParty) string {
|
||||
if len(p.Enrichment) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var meta struct {
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(p.Enrichment, &meta); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
return meta.Status
|
||||
}
|
||||
|
||||
// parseEnrichmentState maps the --state flag to a coredata enrichment
|
||||
// state.
|
||||
func parseEnrichmentState(value string) (coredata.CommonThirdPartyEnrichmentState, error) {
|
||||
switch value {
|
||||
case "queued":
|
||||
return coredata.CommonThirdPartyEnrichmentStateQueued, nil
|
||||
case "enriched":
|
||||
return coredata.CommonThirdPartyEnrichmentStateEnriched, nil
|
||||
case "unenriched":
|
||||
return coredata.CommonThirdPartyEnrichmentStateUnenriched, nil
|
||||
default:
|
||||
return "", fmt.Errorf("invalid --state value %q: valid values are queued, enriched, unenriched", value)
|
||||
}
|
||||
}
|
||||
|
||||
// validEnrichmentStatuses are the run-level statuses the enrichment
|
||||
// worker records in the payload.
|
||||
var validEnrichmentStatuses = map[string]struct{}{
|
||||
"done": {},
|
||||
"partial": {},
|
||||
"failed": {},
|
||||
}
|
||||
|
||||
// resolveCommonThirdPartyID accepts either a common third party GID or a
|
||||
// slug and returns the corresponding id.
|
||||
func resolveCommonThirdPartyID(ctx context.Context, conn pg.Querier, value string) (gid.GID, error) {
|
||||
if id, err := gid.ParseGID(value); err == nil {
|
||||
return id, nil
|
||||
}
|
||||
|
||||
var party coredata.CommonThirdParty
|
||||
if err := party.LoadBySlug(ctx, conn, value); err != nil {
|
||||
if errors.Is(err, coredata.ErrResourceNotFound) {
|
||||
return gid.GID{}, fmt.Errorf("no common third party found for %q (pass a slug or GID)", value)
|
||||
}
|
||||
|
||||
return gid.GID{}, fmt.Errorf("cannot resolve common third party %q: %w", value, err)
|
||||
}
|
||||
|
||||
return party.ID, nil
|
||||
}
|
||||
|
||||
// resolveCommonThirdParty loads a common third party by GID or slug.
|
||||
func resolveCommonThirdParty(ctx context.Context, conn pg.Querier, value string) (coredata.CommonThirdParty, error) {
|
||||
var party coredata.CommonThirdParty
|
||||
|
||||
26
pkg/proboctl/commonthirdparty/list.go
vendored
26
pkg/proboctl/commonthirdparty/list.go
vendored
@@ -31,6 +31,8 @@ func newCmdList(f *cmdutil.Factory) *cobra.Command {
|
||||
flagName string
|
||||
flagCategory string
|
||||
flagKeyword string
|
||||
flagState string
|
||||
flagStatus string
|
||||
flagSort string
|
||||
flagOrder string
|
||||
)
|
||||
@@ -46,6 +48,8 @@ func newCmdList(f *cmdutil.Factory) *cobra.Command {
|
||||
cmd.Flags().StringVar(&flagName, "name", "", "Filter by name substring")
|
||||
cmd.Flags().StringVar(&flagCategory, "category", "", "Filter by category")
|
||||
cmd.Flags().StringVar(&flagKeyword, "keyword", "", "Filter by name/slug substring")
|
||||
cmd.Flags().StringVar(&flagState, "state", "", "Filter by enrichment state (queued, enriched, unenriched)")
|
||||
cmd.Flags().StringVar(&flagStatus, "status", "", "Filter by last enrichment status (done, partial, failed)")
|
||||
cmd.Flags().StringVar(&flagSort, "sort", "name", "Sort field: name, created, updated")
|
||||
cmd.Flags().StringVar(&flagOrder, "order", "", "Sort order: asc, desc (default depends on field)")
|
||||
|
||||
@@ -81,6 +85,23 @@ func newCmdList(f *cmdutil.Factory) *cobra.Command {
|
||||
filter.WithKeyword(&flagKeyword)
|
||||
}
|
||||
|
||||
if flagState != "" {
|
||||
st, err := parseEnrichmentState(flagState)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
filter.WithState(&st)
|
||||
}
|
||||
|
||||
if flagStatus != "" {
|
||||
if _, ok := validEnrichmentStatuses[flagStatus]; !ok {
|
||||
return fmt.Errorf("invalid --status value %q: valid values are done, partial, failed", flagStatus)
|
||||
}
|
||||
|
||||
filter.WithEnrichmentStatus(&flagStatus)
|
||||
}
|
||||
|
||||
pgClient, err := f.PgClient()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -128,14 +149,15 @@ func newCmdList(f *cmdutil.Factory) *cobra.Command {
|
||||
return nil
|
||||
}
|
||||
|
||||
table := clicmdutil.NewTable("ID", "NAME", "SLUG", "CATEGORY", "CREATED", "UPDATED")
|
||||
table := clicmdutil.NewTable("ID", "NAME", "SLUG", "CATEGORY", "STATE", "STATUS", "UPDATED")
|
||||
for _, p := range parties {
|
||||
table.Row(
|
||||
p.ID.String(),
|
||||
p.Name,
|
||||
p.Slug,
|
||||
string(p.Category),
|
||||
p.CreatedAt.Format("2006-01-02 15:04:05"),
|
||||
enrichmentState(p),
|
||||
enrichmentStatus(p),
|
||||
p.UpdatedAt.Format("2006-01-02 15:04:05"),
|
||||
)
|
||||
}
|
||||
|
||||
271
pkg/proboctl/commonthirdparty/reenrich.go
vendored
Normal file
271
pkg/proboctl/commonthirdparty/reenrich.go
vendored
Normal file
@@ -0,0 +1,271 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@probo.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 commonthirdparty
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"go.gearno.de/kit/pg"
|
||||
"go.probo.inc/probo/pkg/coredata"
|
||||
"go.probo.inc/probo/pkg/gid"
|
||||
"go.probo.inc/probo/pkg/proboctl/cmdutil"
|
||||
)
|
||||
|
||||
func newCmdReenrich(f *cmdutil.Factory) *cobra.Command {
|
||||
var (
|
||||
flagIDs []string
|
||||
flagSlugs []string
|
||||
flagCategory string
|
||||
flagKeyword string
|
||||
flagState string
|
||||
flagStatus string
|
||||
flagDryRun bool
|
||||
flagYes bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "reenrich",
|
||||
Short: "Re-enrich common third parties via the enrichment worker",
|
||||
Long: "Re-arm the async common-third-party enrichment worker for selected " +
|
||||
"catalog rows. The worker re-resolves company profile, compliance " +
|
||||
"documents, owned domains, and the logo, merging results over prior " +
|
||||
"per-field provenance so curated seed data and human edits are never " +
|
||||
"overwritten. Already-enriched rows are re-processed. Enrichment is " +
|
||||
"expensive (LLM + browser per row), so a non-empty selection requires " +
|
||||
"--yes; use --dry-run to preview.",
|
||||
Args: cobra.NoArgs,
|
||||
}
|
||||
|
||||
cmd.Flags().StringSliceVar(&flagIDs, "id", nil, "Common third party GID(s) to re-enrich (repeatable)")
|
||||
cmd.Flags().StringSliceVar(&flagSlugs, "slug", nil, "Common third party slug(s) to re-enrich (repeatable)")
|
||||
cmd.Flags().StringVar(&flagCategory, "category", "", "Select rows by category")
|
||||
cmd.Flags().StringVar(&flagKeyword, "keyword", "", "Select rows by a name/slug substring")
|
||||
cmd.Flags().StringVar(&flagState, "state", "", "Select rows by enrichment state (queued, enriched, unenriched)")
|
||||
cmd.Flags().StringVar(&flagStatus, "status", "", "Select rows by last enrichment status (done, partial, failed)")
|
||||
cmd.Flags().BoolVar(&flagDryRun, "dry-run", false, "Print the selected rows without enriching")
|
||||
cmd.Flags().BoolVar(&flagYes, "yes", false, "Skip confirmation")
|
||||
|
||||
cmd.RunE = func(cmd *cobra.Command, args []string) error {
|
||||
ctx := cmd.Context()
|
||||
|
||||
pgClient, err := f.PgClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ids, err := resolveReenrichIDs(
|
||||
ctx,
|
||||
pgClient,
|
||||
flagIDs,
|
||||
flagSlugs,
|
||||
flagCategory,
|
||||
flagKeyword,
|
||||
flagState,
|
||||
flagStatus,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
out := f.IOStreams.Out
|
||||
|
||||
if len(ids) == 0 {
|
||||
_, _ = fmt.Fprintln(out, "No common third parties matched the selection.")
|
||||
return nil
|
||||
}
|
||||
|
||||
if flagDryRun {
|
||||
_, _ = fmt.Fprintf(out, "Would re-enrich %d common third party(ies).\n", len(ids))
|
||||
printSample(out, ids)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
if !flagYes {
|
||||
return fmt.Errorf("about to re-enrich %d common third party(ies); pass --yes to proceed or --dry-run to preview", len(ids))
|
||||
}
|
||||
|
||||
var requeued int64
|
||||
|
||||
if err := pgClient.WithTx(
|
||||
ctx,
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
var parties coredata.CommonThirdParties
|
||||
|
||||
requeued, err = parties.RequestEnrichmentByIDs(ctx, tx, ids)
|
||||
|
||||
return err
|
||||
},
|
||||
); err != nil {
|
||||
return fmt.Errorf("cannot enqueue enrichment: %w", err)
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(out, "Queued %d common third party(ies) for the enrichment worker.\n", requeued)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
// resolveReenrichIDs turns the selection flags into the set of common
|
||||
// third party IDs to re-enrich. Explicit selection (--id and/or --slug)
|
||||
// is used verbatim and the filtering flags do not apply. With no
|
||||
// explicit selection the filtering flags (--category, --keyword,
|
||||
// --state, --status) select across the whole catalog.
|
||||
func resolveReenrichIDs(
|
||||
ctx context.Context,
|
||||
pgClient *pg.Client,
|
||||
rawIDs, rawSlugs []string,
|
||||
category, keyword, state, status string,
|
||||
) ([]gid.GID, error) {
|
||||
explicit := len(rawIDs) > 0 || len(rawSlugs) > 0
|
||||
filtered := category != "" || keyword != "" || state != "" || status != ""
|
||||
|
||||
if explicit && filtered {
|
||||
return nil, fmt.Errorf("--id/--slug cannot be combined with --category, --keyword, --state, or --status")
|
||||
}
|
||||
|
||||
if explicit {
|
||||
return resolveExplicitIDs(ctx, pgClient, rawIDs, rawSlugs)
|
||||
}
|
||||
|
||||
filter, err := buildReenrichFilter(category, keyword, state, status)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var ids []gid.GID
|
||||
|
||||
err = pgClient.WithConn(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
var parties coredata.CommonThirdParties
|
||||
|
||||
ids, err = parties.LoadAllIDs(ctx, conn, filter)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
// resolveExplicitIDs parses the --id GIDs and resolves the --slug values,
|
||||
// preserving order and de-duplicating the combined set.
|
||||
func resolveExplicitIDs(
|
||||
ctx context.Context,
|
||||
pgClient *pg.Client,
|
||||
rawIDs, rawSlugs []string,
|
||||
) ([]gid.GID, error) {
|
||||
seen := make(map[gid.GID]struct{}, len(rawIDs)+len(rawSlugs))
|
||||
ids := make([]gid.GID, 0, len(rawIDs)+len(rawSlugs))
|
||||
|
||||
add := func(id gid.GID) {
|
||||
if _, ok := seen[id]; ok {
|
||||
return
|
||||
}
|
||||
|
||||
seen[id] = struct{}{}
|
||||
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
for _, raw := range rawIDs {
|
||||
id, err := gid.ParseGID(raw)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid --id value %q: %w", raw, err)
|
||||
}
|
||||
|
||||
add(id)
|
||||
}
|
||||
|
||||
if len(rawSlugs) > 0 {
|
||||
if err := pgClient.WithConn(
|
||||
ctx,
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
for _, slug := range rawSlugs {
|
||||
id, err := resolveCommonThirdPartyID(ctx, conn, slug)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
add(id)
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func buildReenrichFilter(category, keyword, state, status string) (*coredata.CommonThirdPartyFilter, error) {
|
||||
filter := coredata.NewCommonThirdPartyFilter(nil)
|
||||
|
||||
if category != "" {
|
||||
cat := coredata.ThirdPartyCategory(category)
|
||||
if !cat.IsValid() {
|
||||
return nil, fmt.Errorf("invalid --category value %q", category)
|
||||
}
|
||||
|
||||
filter.WithCategory(&cat)
|
||||
}
|
||||
|
||||
if keyword != "" {
|
||||
filter.WithKeyword(&keyword)
|
||||
}
|
||||
|
||||
if state != "" {
|
||||
st, err := parseEnrichmentState(state)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filter.WithState(&st)
|
||||
}
|
||||
|
||||
if status != "" {
|
||||
if _, ok := validEnrichmentStatuses[status]; !ok {
|
||||
return nil, fmt.Errorf("invalid --status value %q: valid values are done, partial, failed", status)
|
||||
}
|
||||
|
||||
filter.WithEnrichmentStatus(&status)
|
||||
}
|
||||
|
||||
return filter, nil
|
||||
}
|
||||
|
||||
func printSample(out io.Writer, ids []gid.GID) {
|
||||
const sampleSize = 10
|
||||
|
||||
for i, id := range ids {
|
||||
if i >= sampleSize {
|
||||
_, _ = fmt.Fprintf(out, " ... and %d more\n", len(ids)-sampleSize)
|
||||
break
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(out, " %s\n", id.String())
|
||||
}
|
||||
}
|
||||
102
pkg/proboctl/commonthirdparty/show.go
vendored
102
pkg/proboctl/commonthirdparty/show.go
vendored
@@ -16,8 +16,12 @@ package commonthirdparty
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
"github.com/spf13/cobra"
|
||||
@@ -27,6 +31,31 @@ import (
|
||||
"go.probo.inc/probo/pkg/proboctl/cmdutil"
|
||||
)
|
||||
|
||||
// enrichmentMetadataView mirrors the subset of the enrichment payload
|
||||
// (written by the common-third-party enrichment worker) that show
|
||||
// renders. It is decoded locally to avoid a dependency on the thirdparty
|
||||
// package.
|
||||
type enrichmentMetadataView struct {
|
||||
Model string `json:"model"`
|
||||
AttemptedAt time.Time `json:"attempted_at"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error"`
|
||||
Fields map[string]enrichmentFieldMetaView `json:"fields"`
|
||||
Domains []enrichmentDomainMetaView `json:"domains"`
|
||||
}
|
||||
|
||||
type enrichmentFieldMetaView struct {
|
||||
Confidence float64 `json:"confidence"`
|
||||
SourceURL string `json:"source_url"`
|
||||
Status string `json:"status"`
|
||||
Source string `json:"source"`
|
||||
}
|
||||
|
||||
type enrichmentDomainMetaView struct {
|
||||
Domain string `json:"domain"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
}
|
||||
|
||||
func newCmdShow(f *cmdutil.Factory) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "show <gid|slug>",
|
||||
@@ -115,8 +144,81 @@ func newCmdShow(f *cmdutil.Factory) *cobra.Command {
|
||||
row("Created:", party.CreatedAt.Format("2006-01-02 15:04:05"))
|
||||
row("Updated:", party.UpdatedAt.Format("2006-01-02 15:04:05"))
|
||||
|
||||
row("Enrichment state:", enrichmentState(&party))
|
||||
row("Enrichment attempts:", fmt.Sprintf("%d", party.EnrichmentAttempts))
|
||||
|
||||
if party.EnrichmentRequestedAt != nil {
|
||||
row("Queued at:", party.EnrichmentRequestedAt.Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
|
||||
printEnrichmentDetails(out, label, party)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
// printEnrichmentDetails renders the run-level status and per-field
|
||||
// provenance recorded in the enrichment payload, when present.
|
||||
func printEnrichmentDetails(out io.Writer, label lipgloss.Style, party coredata.CommonThirdParty) {
|
||||
if len(party.Enrichment) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
var meta enrichmentMetadataView
|
||||
if err := json.Unmarshal(party.Enrichment, &meta); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
row := func(name, value string) {
|
||||
_, _ = fmt.Fprintf(out, "%s%s\n", label.Render(name), value)
|
||||
}
|
||||
|
||||
if meta.Status != "" {
|
||||
row("Last run status:", meta.Status)
|
||||
}
|
||||
|
||||
if !meta.AttemptedAt.IsZero() {
|
||||
row("Last attempt:", meta.AttemptedAt.Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
|
||||
if meta.Model != "" {
|
||||
row("Enrichment model:", meta.Model)
|
||||
}
|
||||
|
||||
if meta.Error != "" {
|
||||
row("Last error:", meta.Error)
|
||||
}
|
||||
|
||||
if len(meta.Fields) > 0 {
|
||||
names := make([]string, 0, len(meta.Fields))
|
||||
for name := range meta.Fields {
|
||||
names = append(names, name)
|
||||
}
|
||||
|
||||
sort.Strings(names)
|
||||
|
||||
_, _ = fmt.Fprintln(out)
|
||||
|
||||
table := clicmdutil.NewTable("FIELD", "STATUS", "SOURCE", "CONF")
|
||||
|
||||
for _, name := range names {
|
||||
fm := meta.Fields[name]
|
||||
table.Row(name, fm.Status, fm.Source, fmt.Sprintf("%.2f", fm.Confidence))
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintln(out, table.Render())
|
||||
}
|
||||
|
||||
if len(meta.Domains) > 0 {
|
||||
_, _ = fmt.Fprintln(out)
|
||||
|
||||
table := clicmdutil.NewTable("DISCOVERED DOMAIN", "CONF")
|
||||
for _, d := range meta.Domains {
|
||||
table.Row(d.Domain, fmt.Sprintf("%.2f", d.Confidence))
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintln(out, table.Render())
|
||||
}
|
||||
}
|
||||
|
||||
99
pkg/proboctl/commonthirdparty/stats.go
vendored
Normal file
99
pkg/proboctl/commonthirdparty/stats.go
vendored
Normal file
@@ -0,0 +1,99 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@probo.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 commonthirdparty
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"go.gearno.de/kit/pg"
|
||||
clicmdutil "go.probo.inc/probo/pkg/cmd/cmdutil"
|
||||
"go.probo.inc/probo/pkg/coredata"
|
||||
"go.probo.inc/probo/pkg/proboctl/cmdutil"
|
||||
)
|
||||
|
||||
func newCmdStats(f *cmdutil.Factory) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "stats",
|
||||
Short: "Summarize the common third party catalog by enrichment state and status",
|
||||
Args: cobra.NoArgs,
|
||||
}
|
||||
|
||||
output := clicmdutil.AddOutputFlag(cmd)
|
||||
|
||||
cmd.RunE = func(cmd *cobra.Command, args []string) error {
|
||||
if err := clicmdutil.ValidateOutputFlag(output); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
pgClient, err := f.PgClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stats := map[string]int{}
|
||||
|
||||
if err := pgClient.WithConn(
|
||||
cmd.Context(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
counts := []struct {
|
||||
key string
|
||||
filter *coredata.CommonThirdPartyFilter
|
||||
}{
|
||||
{"total", coredata.NewCommonThirdPartyFilter(nil)},
|
||||
{"queued", coredata.NewCommonThirdPartyFilter(nil).WithState(new(coredata.CommonThirdPartyEnrichmentStateQueued))},
|
||||
{"enriched", coredata.NewCommonThirdPartyFilter(nil).WithState(new(coredata.CommonThirdPartyEnrichmentStateEnriched))},
|
||||
{"unenriched", coredata.NewCommonThirdPartyFilter(nil).WithState(new(coredata.CommonThirdPartyEnrichmentStateUnenriched))},
|
||||
{"status: done", coredata.NewCommonThirdPartyFilter(nil).WithEnrichmentStatus(new("done"))},
|
||||
{"status: partial", coredata.NewCommonThirdPartyFilter(nil).WithEnrichmentStatus(new("partial"))},
|
||||
{"status: failed", coredata.NewCommonThirdPartyFilter(nil).WithEnrichmentStatus(new("failed"))},
|
||||
}
|
||||
|
||||
for _, c := range counts {
|
||||
var parties coredata.CommonThirdParties
|
||||
|
||||
n, err := parties.CountAll(ctx, conn, c.filter)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stats[c.key] = n
|
||||
}
|
||||
|
||||
return nil
|
||||
},
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
order := []string{"total", "queued", "enriched", "unenriched", "status: done", "status: partial", "status: failed"}
|
||||
|
||||
if *output == clicmdutil.OutputJSON {
|
||||
return clicmdutil.PrintJSON(f.IOStreams.Out, stats)
|
||||
}
|
||||
|
||||
table := clicmdutil.NewTable("METRIC", "COUNT")
|
||||
for _, key := range order {
|
||||
table.Row(key, fmt.Sprintf("%d", stats[key]))
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintln(f.IOStreams.Out, table.Render())
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return cmd
|
||||
}
|
||||
@@ -15,6 +15,7 @@
|
||||
package thirdparty
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"go.probo.inc/probo/pkg/uri"
|
||||
@@ -163,13 +164,7 @@ func vendorLabels(name, website string) []string {
|
||||
|
||||
// exactLabelMatch reports whether label equals any vendor label.
|
||||
func exactLabelMatch(label string, vendorLabels []string) bool {
|
||||
for _, vl := range vendorLabels {
|
||||
if label == vl {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
return slices.Contains(vendorLabels, label)
|
||||
}
|
||||
|
||||
// relatedLabelMatch reports whether label is the same as, contains, or is
|
||||
|
||||
Reference in New Issue
Block a user