// +build ignore package main import ( "context" "flag" "fmt" "log" "math/rand" "os" "sync" "sync/atomic" "time" "gorm.io/driver/postgres" "gorm.io/gorm" "gorm.io/gorm/logger" ) // IotCard 简化的卡模型 type IotCard struct { ID uint `gorm:"primaryKey"` ICCID string `gorm:"column:iccid;uniqueIndex:idx_iot_card_iccid,where:deleted_at IS NULL"` CardCategory string `gorm:"column:card_category;default:normal"` CarrierID uint `gorm:"column:carrier_id"` Status int `gorm:"column:status;default:1"` ActivationStatus int `gorm:"column:activation_status;default:0"` RealNameStatus int `gorm:"column:real_name_status;default:0"` NetworkStatus int `gorm:"column:network_status;default:0"` EnablePolling bool `gorm:"column:enable_polling;default:true"` Creator uint `gorm:"column:creator"` Updater uint `gorm:"column:updater"` CreatedAt time.Time UpdatedAt time.Time DeletedAt *time.Time `gorm:"index"` } func (IotCard) TableName() string { return "tb_iot_card" } var ( totalCards = flag.Int("total", 10000000, "要生成的卡数量") batchSize = flag.Int("batch", 10000, "每批插入数量") workers = flag.Int("workers", 10, "并行 worker 数量") startICCID = flag.String("start", "898600000", "起始 ICCID 前缀(9位,总长度不超过20位)") clearOld = flag.Bool("clear", false, "是否清空现有测试卡") insertedCount int64 startTime time.Time ) func main() { flag.Parse() fmt.Println("=== 生成测试卡数据 ===") fmt.Printf("目标数量: %d 张\n", *totalCards) fmt.Printf("批次大小: %d\n", *batchSize) fmt.Printf("并行数: %d\n", *workers) fmt.Println("") // 连接数据库 dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable", os.Getenv("JUNHONG_DATABASE_HOST"), os.Getenv("JUNHONG_DATABASE_PORT"), os.Getenv("JUNHONG_DATABASE_USER"), os.Getenv("JUNHONG_DATABASE_PASSWORD"), os.Getenv("JUNHONG_DATABASE_DBNAME"), ) db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) if err != nil { log.Fatalf("连接数据库失败: %v", err) } // 配置连接池 sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(50) sqlDB.SetMaxIdleConns(25) fmt.Println("✓ 数据库连接成功") // 检查现有卡数量 var existingCount int64 db.Model(&IotCard{}).Count(&existingCount) fmt.Printf("现有卡数量: %d\n", existingCount) if *clearOld { fmt.Println("清空现有测试卡...") // 只删除 ICCID 以 898600000 开头的测试卡 db.Exec("DELETE FROM tb_iot_card WHERE iccid LIKE '898600000%'") fmt.Println("✓ 清空完成") } // 开始生成 startTime = time.Now() ctx := context.Background() // 创建任务通道 taskCh := make(chan int, *workers*2) var wg sync.WaitGroup // 启动 worker for i := 0; i < *workers; i++ { wg.Add(1) go func(workerID int) { defer wg.Done() worker(ctx, db, workerID, taskCh) }(i) } // 分发任务 batches := *totalCards / *batchSize for i := 0; i < batches; i++ { taskCh <- i } close(taskCh) // 等待完成 wg.Wait() elapsed := time.Since(startTime) fmt.Println("") fmt.Println("=== 生成完成 ===") fmt.Printf("总插入: %d 张\n", atomic.LoadInt64(&insertedCount)) fmt.Printf("耗时: %v\n", elapsed) fmt.Printf("速度: %.0f 张/秒\n", float64(atomic.LoadInt64(&insertedCount))/elapsed.Seconds()) // 验证 var finalCount int64 db.Model(&IotCard{}).Count(&finalCount) fmt.Printf("数据库总卡数: %d\n", finalCount) } func worker(ctx context.Context, db *gorm.DB, workerID int, taskCh <-chan int) { rng := rand.New(rand.NewSource(time.Now().UnixNano() + int64(workerID))) for batchIndex := range taskCh { cards := generateBatch(rng, *startICCID, batchIndex, *batchSize) // 批量插入 err := db.WithContext(ctx).CreateInBatches(cards, 1000).Error if err != nil { log.Printf("Worker %d 插入失败: %v", workerID, err) continue } count := atomic.AddInt64(&insertedCount, int64(len(cards))) // 进度报告 if count%100000 == 0 { elapsed := time.Since(startTime).Seconds() speed := float64(count) / elapsed eta := float64(*totalCards-int(count)) / speed fmt.Printf("进度: %d/%d (%.1f%%) | 速度: %.0f/秒 | ETA: %.0f秒\n", count, *totalCards, float64(count)*100/float64(*totalCards), speed, eta) } } } func generateBatch(rng *rand.Rand, iccidPrefix string, batchIndex int, size int) []IotCard { cards := make([]IotCard, size) now := time.Now() for i := 0; i < size; i++ { // 使用前缀 + 序号生成 ICCID(总长度 20 位) // 例如: 898600000 (9位) + 00000000001 (11位) = 20 位 cardIndex := batchIndex*size + i iccid := fmt.Sprintf("%s%011d", iccidPrefix, cardIndex) // 随机分配状态(匹配轮询配置条件) // 实名状态: 0=未实名, 1=实名中, 2=已实名 // 网络状态: 0=停机, 1=正常 // 配置匹配逻辑: // - not_real_name: RealNameStatus == 0 或 1 // - real_name: RealNameStatus == 2 && NetworkStatus != 1 // - activated: RealNameStatus == 2 && NetworkStatus == 1 r := rng.Float64() var realNameStatus, activationStatus, networkStatus int if r < 0.10 { // 10% 未实名 -> 匹配 not_real_name 配置 realNameStatus = 0 activationStatus = 0 networkStatus = 0 } else if r < 0.30 { // 20% 已实名未激活 -> 匹配 real_name 配置 realNameStatus = 2 activationStatus = 0 networkStatus = 0 } else { // 70% 已激活 -> 匹配 activated 配置(流量+套餐检查) realNameStatus = 2 activationStatus = 1 networkStatus = 1 } // 随机卡类型 cardCategory := "normal" if rng.Float64() < 0.05 { cardCategory = "industry" } cards[i] = IotCard{ ICCID: iccid, CardCategory: cardCategory, CarrierID: uint(rng.Intn(3) + 1), // 1-3 运营商 Status: 1, ActivationStatus: activationStatus, RealNameStatus: realNameStatus, NetworkStatus: networkStatus, EnablePolling: true, Creator: 1, Updater: 1, CreatedAt: now, UpdatedAt: now, } } return cards }