129 lines
3.5 KiB
Go
129 lines
3.5 KiB
Go
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)")
|
|
|
|
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); 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) 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
|
|
} |