Files
kefu/im/backend/internal/testusers/seed.go
T
2026-09-03 08:38:17 +08:00

112 lines
4.3 KiB
Go

package testusers
import (
"context"
"database/sql"
"fmt"
)
type Result struct {
Batch string `json:"batch"`
Created int `json:"created"`
Skipped int `json:"skipped"`
Male int `json:"male"`
Female int `json:"female"`
IDs []int64 `json:"ids"`
}
// Seed is atomic and repeatable: rerunning a complete batch is a no-op.
// Conflicting real users and incomplete/modified batches are never overwritten.
func Seed(ctx context.Context, db *sql.DB, publicBase, confirmDatabase string) (Result, error) {
result := Result{Batch: Batch, IDs: []int64{}}
profiles, err := Generate(publicBase)
if err != nil {
return result, err
}
conn, err := db.Conn(ctx)
if err != nil {
return result, err
}
defer conn.Close()
var database string
if err = conn.QueryRowContext(ctx, `SELECT DATABASE()`).Scan(&database); err != nil {
return result, err
}
if confirmDatabase == "" || database != confirmDatabase {
return result, fmt.Errorf("database confirmation does not match the connected database")
}
var locked int
if err = conn.QueryRowContext(ctx, `SELECT GET_LOCK(?,10)`, "xingyu:testusers:"+Batch).Scan(&locked); err != nil || locked != 1 {
return result, fmt.Errorf("could not acquire test-user batch lock")
}
defer func() { _, _ = conn.ExecContext(context.Background(), `DO RELEASE_LOCK(?)`, "xingyu:testusers:"+Batch) }()
tx, err := conn.BeginTx(ctx, nil)
if err != nil {
return result, err
}
defer tx.Rollback()
var batchCount int
if err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM users WHERE test_batch=?`, Batch).Scan(&batchCount); err != nil {
return result, fmt.Errorf("test-user schema unavailable; apply migration 028 first: %w", err)
}
if batchCount != 0 && batchCount != Count {
return result, fmt.Errorf("batch has %d users, expected 0 or %d; refusing to modify partial data", batchCount, Count)
}
for _, p := range profiles {
var id int64
var isTest bool
var existingBatch string
var deleted sql.NullTime
err = tx.QueryRowContext(ctx, `SELECT id,is_test,test_batch,deleted_at FROM users WHERE public_id=? FOR UPDATE`, p.PublicID).Scan(&id, &isTest, &existingBatch, &deleted)
if err == nil {
if batchCount != Count || !isTest || existingBatch != Batch || deleted.Valid {
return Result{}, fmt.Errorf("public ID %s conflicts with an existing or deleted user; no users changed", p.PublicID)
}
var gender int
if err = tx.QueryRowContext(ctx, `SELECT gender FROM user_profiles WHERE user_id=?`, id).Scan(&gender); err != nil || gender != p.Gender {
return Result{}, fmt.Errorf("existing test profile %s has been modified or is incomplete", p.PublicID)
}
result.Skipped++
} else if err == sql.ErrNoRows && batchCount == 0 {
res, insertErr := tx.ExecContext(ctx, `INSERT INTO users (public_id,password_hash,is_test,test_batch,status) VALUES (?,'!TEST_PROFILE_NO_LOGIN',1,?,1)`, p.PublicID, Batch)
if insertErr != nil {
return Result{}, insertErr
}
id, err = res.LastInsertId()
if err != nil {
return Result{}, err
}
_, err = tx.ExecContext(ctx, `INSERT INTO user_profiles (user_id,nickname,avatar_url,cover_url,gender,birthday,height_cm,city_code,city_name,occupation,bio,profile_score) VALUES (?,?,?,'',?,?,?,?,?,?,?,80)`, id, p.Nickname, p.Avatar, p.Gender, p.Birthday, p.Height, p.CityCode, p.City, p.Occupation, p.Bio)
if err != nil {
return Result{}, err
}
// Fixture coordinates stay close to the declared city center and are
// explicitly marked as imported, never as a real GPS observation.
_, err = tx.ExecContext(ctx, `INSERT INTO user_privacy_settings (user_id,distance_visible,online_visible,last_active_visible,allow_profile_visit_record) VALUES (?,1,1,0,0)`, id)
if err != nil {
return Result{}, err
}
_, err = tx.ExecContext(ctx, `INSERT INTO user_location_states(user_id,city_code,location_cell,latitude,longitude,source) VALUES(?,?,'fixture',?,?,'fixture')`, id, p.CityCode, p.Latitude, p.Longitude)
if err != nil {
return Result{}, err
}
result.Created++
} else {
if err == sql.ErrNoRows {
err = fmt.Errorf("batch identity mismatch at %s", p.PublicID)
}
return Result{}, err
}
if p.Gender == 1 {
result.Male++
} else {
result.Female++
}
result.IDs = append(result.IDs, id)
}
if err = tx.Commit(); err != nil {
return Result{}, err
}
return result, nil
}