GPT_Management/backend/internal/db/migrations.go

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);
`