diff --git a/backend/scripts/loadgen/seed/main.go b/backend/scripts/loadgen/seed/main.go new file mode 100644 index 0000000..81cddbb --- /dev/null +++ b/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 +} \ No newline at end of file