// Copyright (c) 2026 Probo Inc . // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal // in the Software without restriction, including without limitation the rights // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell // copies of the Software, and to permit persons to whom the Software is // furnished to do so, subject to the following conditions: // // The above copyright notice and this permission notice shall be included in // all copies or substantial portions of the Software. // // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // 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()) } }