feat(seed): main.go skeleton with config, DB open, cleanup subcommand

backend/scripts/loadgen/seed/main.go

@@ -0,0 +1,130 @@
+package main
+
+import (
+	"database/sql"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+
+	_ "github.com/lib/pq"
+)
+
+const (
+	LoadtestStarID  = int64(999900)
+	LoadtestUserMin = int64(30000001)
+	LoadtestUserMax = int64(30001000)
+)
+
+type Config struct {
+	JWTSecret string
+	DBHost    string
+	DBPort    int
+	DBName    string
+	DBUser    string
+	DBPass    string
+	Reset     bool
+	ResetTok  bool
+}
+
+func main() {
+	cleanup := flag.Bool("cleanup", false, "run cleanup (default: keep baseline)")
+	cleanupFull := flag.Bool("full", false, "with -cleanup: also delete users/stars")
+	cleanupStarID := flag.Int64("cleanup-star-id", LoadtestStarID, "star_id to clean (safety)")
+	flag.Parse()
+
+	cfg := parseFlags()
+
+	if *cleanup {
+		db, err := openDB(cfg)
+		if err != nil {
+			log.Fatalf("open db: %v", err)
+		}
+		defer db.Close()
+		if err := Cleanup(db, *cleanupStarID, *cleanupFull); err != nil {
+			log.Fatalf("cleanup: %v", err)
+		}
+		log.Println("✅ cleanup done")
+		return
+	}
+
+	if cfg.ResetTok {
+		if err := GenerateTokensForLoadtest(cfg); err != nil {
+			log.Fatalf("generate tokens: %v", err)
+		}
+		return
+	}
+
+	db, err := openDB(cfg)
+	if err != nil {
+		log.Fatalf("open db: %v", err)
+	}
+	defer db.Close()
+
+	if cfg.Reset {
+		log.Println("⚠️ --reset enabled, deleting existing loadtest data first")
+		if err := Cleanup(db, LoadtestStarID, false); err != nil {
+			log.Fatalf("reset: %v", err)
+		}
+	}
+
+	if err := runSeed(db, cfg); err != nil {
+		log.Fatalf("seed failed: %v", err)
+	}
+	if err := GenerateTokensForLoadtest(cfg); err != nil {
+		log.Fatalf("generate tokens: %v", err)
+	}
+	log.Println("✅ seed + tokens completed")
+}
+
+func parseFlags() *Config {
+	cfg := &Config{}
+	flag.StringVar(&cfg.JWTSecret, "jwt-secret", os.Getenv("JWT_SECRET"), "JWT secret (或 $JWT_SECRET)")
+	flag.StringVar(&cfg.DBHost, "db-host", "localhost", "PG host")
+	flag.IntVar(&cfg.DBPort, "db-port", 5432, "PG port")
+	flag.StringVar(&cfg.DBName, "db-name", "topfans", "PG database name (本地为 'top-fans' 带横线)")
+	flag.StringVar(&cfg.DBUser, "db-user", "postgres", "PG user")
+	flag.StringVar(&cfg.DBPass, "db-password", os.Getenv("DB_PASSWORD"), "PG password (或 $DB_PASSWORD)")
+	flag.BoolVar(&cfg.Reset, "reset", false, "delete existing test data before seed")
+	flag.BoolVar(&cfg.ResetTok, "reset-tokens", false, "only re-sign tokens, don't touch data")
+	flag.Parse()
+	return cfg
+}
+
+func openDB(cfg *Config) (*sql.DB, error) {
+	dsn := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+		cfg.DBHost, cfg.DBPort, cfg.DBUser, cfg.DBPass, cfg.DBName)
+	return sql.Open("postgres", dsn)
+}
+
+func runSeed(db *sql.DB, cfg *Config) error {
+	if err := SeedStars(db); err != nil {
+		return fmt.Errorf("seed stars: %w", err)
+	}
+	log.Println("✓ stars seeded")
+	if err := SeedUsers(db); err != nil {
+		return fmt.Errorf("seed users: %w", err)
+	}
+	log.Println("✓ 1000 users seeded")
+	if err := SeedProfiles(db); err != nil {
+		return fmt.Errorf("seed profiles: %w", err)
+	}
+	log.Println("✓ 1000 fan_profiles + crystal seeded")
+	if err := SeedAssets(db); err != nil {
+		return fmt.Errorf("seed assets: %w", err)
+	}
+	log.Println("✓ 5000 assets seeded")
+	if err := SeedSlotsAndExhibits(db); err != nil {
+		return fmt.Errorf("seed slots+exhibits: %w", err)
+	}
+	log.Println("✓ 3000 booth_slots + 2000 exhibitions seeded")
+	if err := SeedFriendships(db); err != nil {
+		return fmt.Errorf("seed friendships: %w", err)
+	}
+	log.Println("✓ 10000 friendships seeded")
+	if err := ResetSequences(db); err != nil {
+		return fmt.Errorf("reset sequences: %w", err)
+	}
+	log.Println("✓ sequences reset")
+	return nil
+}