package db import ( "database/sql" "fmt" "log" "os" "time" "gpt-manager-go/internal/auth" ) // Migrate 执行数据库迁移 func Migrate(db *sql.DB) error { log.Println("Starting database migration...") // 创建表的 SQL 语句 migrations := []string{ createAdminsTable, createChatGPTAccountsTable, createCardKeysTable, createInvitationsTable, createAPIKeysTable, } for i, migration := range migrations { if _, err := db.Exec(migration); err != nil { return fmt.Errorf("migration %d failed: %w", i+1, err) } } // 创建默认管理员 if err := CreateDefaultAdmin(db); err != nil { log.Printf("Warning: Failed to create default admin: %v", err) } log.Println("Database migration completed successfully") return nil } // CreateDefaultAdmin 创建默认管理员(如果不存在) func CreateDefaultAdmin(db *sql.DB) error { username := os.Getenv("ADMIN_USERNAME") email := os.Getenv("ADMIN_EMAIL") password := os.Getenv("ADMIN_PASSWORD") // 如果没有配置管理员信息,跳过 if username == "" || email == "" || password == "" { log.Println("No default admin configuration found, skipping...") return nil } // 检查是否已存在管理员 var count int err := db.QueryRow("SELECT COUNT(*) FROM admins").Scan(&count) if err != nil { return fmt.Errorf("failed to check admins: %w", err) } if count > 0 { log.Println("Admin already exists, skipping default admin creation...") return nil } // 创建密码哈希 passwordHash, err := auth.HashPassword(password) if err != nil { return fmt.Errorf("failed to hash password: %w", err) } // 创建管理员 _, err = db.Exec(` INSERT INTO admins (username, email, password_hash, is_super_admin, is_active, created_at) VALUES ($1, $2, $3, $4, $5, $6) `, username, email, passwordHash, true, true, time.Now()) if err != nil { return fmt.Errorf("failed to create admin: %w", err) } log.Printf("Default admin '%s' created successfully", username) return nil } const createAdminsTable = ` CREATE TABLE IF NOT EXISTS admins ( id SERIAL PRIMARY KEY, username VARCHAR(50) NOT NULL UNIQUE, email VARCHAR(120) NOT NULL UNIQUE, password_hash VARCHAR(128) NOT NULL, is_super_admin BOOLEAN NOT NULL DEFAULT FALSE, is_active BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, last_login TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_admins_username ON admins(username); CREATE INDEX IF NOT EXISTS idx_admins_email ON admins(email); ` const createChatGPTAccountsTable = ` CREATE TABLE IF NOT EXISTS chatgpt_accounts ( id SERIAL PRIMARY KEY, name VARCHAR(100) NOT NULL, auth_token TEXT NOT NULL, team_account_id VARCHAR(100) NOT NULL UNIQUE, seats_in_use INTEGER NOT NULL DEFAULT 0, seats_entitled INTEGER NOT NULL DEFAULT 0, active_start TIMESTAMP, active_until TIMESTAMP, is_active BOOLEAN NOT NULL DEFAULT TRUE, consecutive_failures INTEGER NOT NULL DEFAULT 0, last_check TIMESTAMP, last_used TIMESTAMP, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_chatgpt_accounts_is_active ON chatgpt_accounts(is_active); CREATE INDEX IF NOT EXISTS idx_chatgpt_accounts_team_account_id ON chatgpt_accounts(team_account_id); ` const createCardKeysTable = ` CREATE TABLE IF NOT EXISTS card_keys ( id SERIAL PRIMARY KEY, key VARCHAR(19) NOT NULL UNIQUE, max_uses INTEGER NOT NULL DEFAULT 1, used_count INTEGER NOT NULL DEFAULT 0, validity_type VARCHAR(20) NOT NULL, expires_at TIMESTAMP NOT NULL, is_active BOOLEAN NOT NULL DEFAULT TRUE, created_by_id INTEGER NOT NULL REFERENCES admins(id), created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_card_keys_key ON card_keys(key); CREATE INDEX IF NOT EXISTS idx_card_keys_expires_at ON card_keys(expires_at); CREATE INDEX IF NOT EXISTS idx_card_keys_is_active ON card_keys(is_active); ` const createInvitationsTable = ` CREATE TABLE IF NOT EXISTS invitations ( id SERIAL PRIMARY KEY, card_key_id INTEGER REFERENCES card_keys(id), account_id INTEGER NOT NULL REFERENCES chatgpt_accounts(id), invited_email VARCHAR(120) NOT NULL, status VARCHAR(20) NOT NULL DEFAULT 'pending', error_message TEXT, expires_at TIMESTAMP, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_invitations_card_key_id ON invitations(card_key_id); CREATE INDEX IF NOT EXISTS idx_invitations_account_id ON invitations(account_id); CREATE INDEX IF NOT EXISTS idx_invitations_invited_email ON invitations(invited_email); CREATE INDEX IF NOT EXISTS idx_invitations_status ON invitations(status); CREATE INDEX IF NOT EXISTS idx_invitations_created_at ON invitations(created_at); ` const createAPIKeysTable = ` CREATE TABLE IF NOT EXISTS api_keys ( id SERIAL PRIMARY KEY, key VARCHAR(64) NOT NULL UNIQUE, name VARCHAR(100) NOT NULL, created_by_id INTEGER NOT NULL REFERENCES admins(id), is_active BOOLEAN NOT NULL DEFAULT TRUE, rate_limit INTEGER NOT NULL DEFAULT 60, allowed_ips TEXT DEFAULT '[]', last_used TIMESTAMP, request_count INTEGER NOT NULL DEFAULT 0, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_api_keys_key ON api_keys(key); CREATE INDEX IF NOT EXISTS idx_api_keys_is_active ON api_keys(is_active); `