sqlx_migrate.go

package outbox

import (
	"context"
	"fmt"
	"strings"

	"github.com/jmoiron/sqlx"
)

// AutoMigrateSqlx 自动创建或更新表结构 (用于 Sqlx)
func AutoMigrateSqlx(db *sqlx.DB) error {
	ctx := context.Background()

	// 获取数据库类型
	driverName := db.DriverName()

	// 创建或更新 Published 表
	if err := createOrUpdatePublishedTable(ctx, db, driverName); err != nil {
		return fmt.Errorf("create or update published table error: %w", err)
	}

	// 创建或更新 Received 表
	if err := createOrUpdateReceivedTable(ctx, db, driverName); err != nil {
		return fmt.Errorf("create or update received table error: %w", err)
	}

	return nil
}

// createOrUpdatePublishedTable 创建或更新 Published 表
func createOrUpdatePublishedTable(ctx context.Context, db *sqlx.DB, driverName string) error {
	tableName := PublishedTableName

	// 检查表是否存在
	exists, err := tableExists(ctx, db, tableName, driverName)
	if err != nil {
		return err
	}

	if !exists {
		// 表不存在，创建新表
		return createPublishedTable(ctx, db, driverName)
	}

	// 表已存在，检查并添加缺失的列
	return alterPublishedTable(ctx, db, tableName, driverName)
}

// createOrUpdateReceivedTable 创建或更新 Received 表
func createOrUpdateReceivedTable(ctx context.Context, db *sqlx.DB, driverName string) error {
	tableName := ReceivedTableName

	// 检查表是否存在
	exists, err := tableExists(ctx, db, tableName, driverName)
	if err != nil {
		return err
	}

	if !exists {
		// 表不存在，创建新表
		return createReceivedTable(ctx, db, driverName)
	}

	// 表已存在，检查并添加缺失的列
	return alterReceivedTable(ctx, db, tableName, driverName)
}

// tableExists 检查表是否存在
func tableExists(ctx context.Context, db *sqlx.DB, tableName, driverName string) (bool, error) {
	var query string
	var count int

	switch driverName {
	case "mysql":
		query = `SELECT COUNT(*) FROM information_schema.tables 
				 WHERE table_schema = DATABASE() AND table_name = ?`
	case "postgres":
		query = `SELECT COUNT(*) FROM information_schema.tables 
				 WHERE table_schema = 'public' AND table_name = $1`
	case "sqlite3":
		query = `SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?`
	default:
		return false, fmt.Errorf("unsupported database driver: %s", driverName)
	}

	err := db.GetContext(ctx, &count, query, tableName)
	if err != nil {
		return false, err
	}

	return count > 0, nil
}

// columnExists 检查列是否存在
func columnExists(ctx context.Context, db *sqlx.DB, tableName, columnName, driverName string) (bool, error) {
	var query string
	var count int

	switch driverName {
	case "mysql":
		query = `SELECT COUNT(*) FROM information_schema.columns 
				 WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`
	case "postgres":
		query = `SELECT COUNT(*) FROM information_schema.columns 
				 WHERE table_schema = 'public' AND table_name = $1 AND column_name = $2`
	case "sqlite3":
		// SQLite 需要特殊处理
		rows, err := db.QueryxContext(ctx, fmt.Sprintf("PRAGMA table_info(%s)", tableName))
		if err != nil {
			return false, err
		}
		defer rows.Close()

		for rows.Next() {
			row := make(map[string]interface{})
			if err := rows.MapScan(row); err != nil {
				return false, err
			}
			if name, ok := row["name"].(string); ok && name == columnName {
				return true, nil
			}
		}
		return false, nil
	default:
		return false, fmt.Errorf("unsupported database driver: %s", driverName)
	}

	if driverName != "sqlite3" {
		err := db.GetContext(ctx, &count, query, tableName, columnName)
		if err != nil {
			return false, err
		}
		return count > 0, nil
	}

	return false, nil
}

// createPublishedTable 创建 Published 表
func createPublishedTable(ctx context.Context, db *sqlx.DB, driverName string) error {
	var createSQL string

	switch driverName {
	case "mysql":
		createSQL = `CREATE TABLE ` + PublishedTableName + ` (
			id BIGINT(20) PRIMARY KEY,
			version VARCHAR(20),
			topic VARCHAR(200) NOT NULL,
			value LONGTEXT,
			retries INT(11) DEFAULT 0,
			created_at DATETIME NOT NULL,
			expires_at DATETIME,
			status_name VARCHAR(40) NOT NULL,
			INDEX idx_status_name (status_name),
			INDEX idx_version_status (version, status_name)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`
	case "postgres":
		createSQL = `CREATE TABLE ` + PublishedTableName + ` (
			id BIGINT PRIMARY KEY,
			version VARCHAR(20),
			topic VARCHAR(200) NOT NULL,
			value TEXT,
			retries INTEGER DEFAULT 0,
			created_at TIMESTAMP NOT NULL,
			expires_at TIMESTAMP,
			status_name VARCHAR(40) NOT NULL
		)`
	case "sqlite3":
		createSQL = `CREATE TABLE ` + PublishedTableName + ` (
			id INTEGER PRIMARY KEY,
			version TEXT,
			topic TEXT NOT NULL,
			value TEXT,
			retries INTEGER DEFAULT 0,
			created_at DATETIME NOT NULL,
			expires_at DATETIME,
			status_name TEXT NOT NULL
		)`
	default:
		return fmt.Errorf("unsupported database driver: %s", driverName)
	}

	_, err := db.ExecContext(ctx, createSQL)
	if err != nil {
		return err
	}

	// 为 PostgreSQL 创建索引
	if driverName == "postgres" {
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_status_name ON `+PublishedTableName+` (status_name)`)
		if err != nil {
			return err
		}
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_version_status ON `+PublishedTableName+` (version, status_name)`)
		if err != nil {
			return err
		}
	}

	// 为 SQLite 创建索引
	if driverName == "sqlite3" {
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_status_name ON `+PublishedTableName+` (status_name)`)
		if err != nil {
			return err
		}
	}

	return nil
}

// createReceivedTable 创建 Received 表
func createReceivedTable(ctx context.Context, db *sqlx.DB, driverName string) error {
	var createSQL string

	switch driverName {
	case "mysql":
		createSQL = `CREATE TABLE ` + ReceivedTableName + ` (
			id BIGINT(20) PRIMARY KEY,
			version VARCHAR(20),
			name VARCHAR(200) NOT NULL,
			` + "`group`" + ` VARCHAR(200) NOT NULL,
			value LONGTEXT,
			retries INT(11) DEFAULT 0,
			created_at DATETIME NOT NULL,
			expires_at DATETIME,
			status_name VARCHAR(40) NOT NULL,
			INDEX idx_status_name (status_name),
			INDEX idx_version_status (version, status_name)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`
	case "postgres":
		createSQL = `CREATE TABLE ` + ReceivedTableName + ` (
			id BIGINT PRIMARY KEY,
			version VARCHAR(20),
			name VARCHAR(200) NOT NULL,
			"group" VARCHAR(200) NOT NULL,
			value TEXT,
			retries INTEGER DEFAULT 0,
			created_at TIMESTAMP NOT NULL,
			expires_at TIMESTAMP,
			status_name VARCHAR(40) NOT NULL
		)`
	case "sqlite3":
		createSQL = `CREATE TABLE ` + ReceivedTableName + ` (
			id INTEGER PRIMARY KEY,
			version TEXT,
			name TEXT NOT NULL,
			"group" TEXT NOT NULL,
			value TEXT,
			retries INTEGER DEFAULT 0,
			created_at DATETIME NOT NULL,
			expires_at DATETIME,
			status_name TEXT NOT NULL
		)`
	default:
		return fmt.Errorf("unsupported database driver: %s", driverName)
	}

	_, err := db.ExecContext(ctx, createSQL)
	if err != nil {
		return err
	}

	// 为 PostgreSQL 创建索引
	if driverName == "postgres" {
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_received_status_name ON `+ReceivedTableName+` (status_name)`)
		if err != nil {
			return err
		}
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_received_version_status ON `+ReceivedTableName+` (version, status_name)`)
		if err != nil {
			return err
		}
	}

	// 为 SQLite 创建索引
	if driverName == "sqlite3" {
		_, err = db.ExecContext(ctx, `CREATE INDEX idx_received_status_name ON `+ReceivedTableName+` (status_name)`)
		if err != nil {
			return err
		}
	}

	return nil
}

// alterPublishedTable 调整 Published 表结构
func alterPublishedTable(ctx context.Context, db *sqlx.DB, tableName, driverName string) error {
	// 定义所有必需的列
	requiredColumns := map[string]string{
		"id":          getColumnType(driverName, "id"),
		"version":     getColumnType(driverName, "version"),
		"topic":       getColumnType(driverName, "topic"),
		"value":       getColumnType(driverName, "value"),
		"retries":     getColumnType(driverName, "retries"),
		"created_at":  getColumnType(driverName, "created_at"),
		"expires_at":  getColumnType(driverName, "expires_at"),
		"status_name": getColumnType(driverName, "status_name"),
	}

	// 检查并添加缺失的列
	for columnName, columnType := range requiredColumns {
		exists, err := columnExists(ctx, db, tableName, columnName, driverName)
		if err != nil {
			return err
		}

		if !exists {
			if err := addColumn(ctx, db, tableName, columnName, columnType, driverName); err != nil {
				return fmt.Errorf("add column %s error: %w", columnName, err)
			}
		}
	}

	return nil
}

// alterReceivedTable 调整 Received 表结构
func alterReceivedTable(ctx context.Context, db *sqlx.DB, tableName, driverName string) error {
	// 定义所有必需的列
	requiredColumns := map[string]string{
		"id":          getColumnType(driverName, "id"),
		"version":     getColumnType(driverName, "version"),
		"name":        getColumnType(driverName, "name"),
		"group":       getColumnType(driverName, "group"),
		"value":       getColumnType(driverName, "value"),
		"retries":     getColumnType(driverName, "retries"),
		"created_at":  getColumnType(driverName, "created_at"),
		"expires_at":  getColumnType(driverName, "expires_at"),
		"status_name": getColumnType(driverName, "status_name"),
	}

	// 检查并添加缺失的列
	for columnName, columnType := range requiredColumns {
		exists, err := columnExists(ctx, db, tableName, columnName, driverName)
		if err != nil {
			return err
		}

		if !exists {
			if err := addColumn(ctx, db, tableName, columnName, columnType, driverName); err != nil {
				return fmt.Errorf("add column %s error: %w", columnName, err)
			}
		}
	}

	return nil
}

// addColumn 添加列
func addColumn(ctx context.Context, db *sqlx.DB, tableName, columnName, columnType, driverName string) error {
	var alterSQL string

	// 处理保留字
	escapedColumnName := columnName
	if columnName == "group" {
		switch driverName {
		case "mysql":
			escapedColumnName = "`group`"
		case "postgres":
			escapedColumnName = `"group"`
		}
	}

	switch driverName {
	case "mysql":
		alterSQL = fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, escapedColumnName, columnType)
	case "postgres":
		alterSQL = fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, escapedColumnName, columnType)
	case "sqlite3":
		alterSQL = fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", tableName, escapedColumnName, columnType)
	default:
		return fmt.Errorf("unsupported database driver: %s", driverName)
	}

	_, err := db.ExecContext(ctx, alterSQL)
	return err
}

// getColumnType 根据数据库类型获取列类型定义
func getColumnType(driverName, columnName string) string {
	switch driverName {
	case "mysql":
		switch columnName {
		case "id":
			return "BIGINT(20) PRIMARY KEY"
		case "version":
			return "VARCHAR(20)"
		case "topic", "name", "group":
			return "VARCHAR(200) NOT NULL"
		case "value":
			return "LONGTEXT"
		case "retries":
			return "INT(11) DEFAULT 0"
		case "created_at":
			return "DATETIME NOT NULL"
		case "expires_at":
			return "DATETIME"
		case "status_name":
			return "VARCHAR(40) NOT NULL"
		}
	case "postgres":
		switch columnName {
		case "id":
			return "BIGINT PRIMARY KEY"
		case "version":
			return "VARCHAR(20)"
		case "topic", "name", "group":
			return "VARCHAR(200) NOT NULL"
		case "value":
			return "TEXT"
		case "retries":
			return "INTEGER DEFAULT 0"
		case "created_at":
			return "TIMESTAMP NOT NULL"
		case "expires_at":
			return "TIMESTAMP"
		case "status_name":
			return "VARCHAR(40) NOT NULL"
		}
	case "sqlite3":
		switch columnName {
		case "id":
			return "INTEGER PRIMARY KEY"
		case "version":
			return "TEXT"
		case "topic", "name", "group":
			return "TEXT NOT NULL"
		case "value":
			return "TEXT"
		case "retries":
			return "INTEGER DEFAULT 0"
		case "created_at":
			return "DATETIME NOT NULL"
		case "expires_at":
			return "DATETIME"
		case "status_name":
			return "TEXT NOT NULL"
		}
	}
	return ""
}

// GetDriverName 获取数据库驱动名称（辅助函数）
func GetDriverName(db *sqlx.DB) string {
	name := db.DriverName()
	// 标准化驱动名称
	name = strings.ToLower(name)
	if strings.Contains(name, "mysql") {
		return "mysql"
	}
	if strings.Contains(name, "postgres") || strings.Contains(name, "pgx") {
		return "postgres"
	}
	if strings.Contains(name, "sqlite") {
		return "sqlite3"
	}
	return name
}