descrybe/apps/api/cmd/migrator/company_plans_repair.go

package main

import (
	"context"
	"fmt"
	"log"
	"strings"

	"github.com/descrybe/descrybe-v2/apps/api/internal/billing"
	"github.com/jackc/pgx/v5/pgxpool"
)

// runCompaniesWithoutPlansRepair lists and/or assigns plans for companies that
// have no active company_plans row. Postgres-only; never deletes or overwrites
// an existing active plan. Live writes require confirm=true (no blind assigns).
func runCompaniesWithoutPlansRepair(postgresURL, planName string, listOnly, assign bool, dryRun, confirm bool) {
	if postgresURL == "" {
		log.Fatal("-postgres / DATABASE_URL is required for companies-without-plans tooling")
	}
	planName, err := normalizeAssignPlanName(planName, assign)
	if err != nil {
		log.Fatal(err)
	}
	if !listOnly && !assign {
		log.Fatal("pass -list-companies-without-plans and/or -assign-missing-plans")
	}
	if err := guardLiveMutation(assign, dryRun, confirm, "-assign-missing-plans"); err != nil {
		log.Fatal(err)
	}

	ctx := context.Background()
	pg, err := pgxpool.New(ctx, postgresURL)
	if err != nil {
		log.Fatalf("postgres: %v", err)
	}
	defer pg.Close()

	svc := &billing.Service{Pool: pg}
	if err := svc.EnsureDefaultPlans(ctx); err != nil {
		log.Fatalf("ensure default plans: %v", err)
	}

	total, err := svc.CountCompaniesWithoutActivePlan(ctx)
	if err != nil {
		log.Fatalf("count companies without active plan: %v", err)
	}
	fmt.Printf("companies_without_active_plan: %d\n", total)

	const page = 200
	offset := 0
	listed := 0
	assigned := 0
	skipped := 0
	a1Skipped := 0

	for {
		rows, err := svc.ListCompaniesWithoutActivePlan(ctx, page, offset)
		if err != nil {
			log.Fatalf("list companies without active plan: %v", err)
		}
		if len(rows) == 0 {
			break
		}
		pageAssigned := 0
		for _, c := range rows {
			listed++
			a1 := billing.IsA1CohortCompany(c.LegacyCompanyID, c.Name)
			if listOnly || !assign {
				fmt.Printf("  %s\t%s\t%s\ta1=%v\n", c.ID, c.Name, c.Language, a1)
			}
			mutate, skipReason := decideAssignMissingPlan(assign, c.LegacyCompanyID, c.Name)
			if !mutate {
				if assign && skipReason != "" {
					fmt.Printf("skip\t%s\t%s\t%s\n", c.ID, c.Name, skipReason)
					skipped++
					a1Skipped++
				}
				continue
			}
			if dryRun {
				fmt.Printf("dry-run: would assign plan %q to company %s (%s)\n", planName, c.ID, c.Name)
				assigned++
				continue
			}
			ok, err := svc.AssignPlanByNameIfMissing(ctx, c.ID, planName)
			if err != nil {
				log.Printf("assign plan %q to company %s: %v", planName, c.ID, err)
				skipped++
				continue
			}
			if !ok {
				skipped++
				continue
			}
			fmt.Printf("assigned plan %q to company %s (%s)\n", planName, c.ID, c.Name)
			assigned++
			pageAssigned++
		}
		if len(rows) < page {
			break
		}
		if assign && !dryRun {
			// Live assigns shrink the result set, so restart from offset 0.
			// If this page assigned nothing (e.g. all A1 skips), advance offset
			// so we cannot spin forever on the same unassignable rows.
			if pageAssigned == 0 {
				offset += page
			} else {
				offset = 0
			}
			continue
		}
		offset += page
	}

	if assign {
		fmt.Printf("listed=%d assigned=%d skipped=%d a1_skipped=%d dry_run=%v plan=%q\n", listed, assigned, skipped, a1Skipped, dryRun, planName)
	} else {
		fmt.Printf("listed=%d\n", listed)
	}
}

// decideAssignMissingPlan is the assign gate used by dry-run and -confirm.
// A1 cohort companies always skip (never get Free/other fallback), even when confirm=true.
func decideAssignMissingPlan(assign bool, legacyCompanyID, companyName string) (mutate bool, skipReason string) {
	if !assign {
		return false, ""
	}
	if billing.IsA1CohortCompany(legacyCompanyID, companyName) {
		return false, "a1_cohort"
	}
	return true, ""
}

func normalizeAssignPlanName(planName string, assign bool) (string, error) {
	planName = strings.TrimSpace(planName)
	if assign && planName == "" {
		return "", fmt.Errorf("-plan-name is required with -assign-missing-plans (e.g. Free)")
	}
	return planName, nil
}

// guardLiveAssign is kept for tests; prefer guardLiveMutation for new call sites.
func guardLiveAssign(assign, dryRun, confirm bool) error {
	return guardLiveMutation(assign, dryRun, confirm, "-assign-missing-plans")
}

func resolveFallbackPlanID(ctx context.Context, pg *pgxpool.Pool, planName string) (int64, error) {
	planName = strings.TrimSpace(planName)
	if planName == "" {
		return 0, nil
	}
	svc := &billing.Service{Pool: pg}
	if err := svc.EnsureDefaultPlans(ctx); err != nil {
		return 0, err
	}
	return svc.PlanIDByName(ctx, planName)
}