golib/jwt/tablemanager.go

package jwt

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/pkg/errors"
)

// TableManager handles table creation, existence checks, and cleanup configuration.
type TableManager struct {
	dbType      DatabaseType
	tableConfig TableConfig
	db          *sql.DB
}

// NewTableManager creates a new TableManager instance.
func NewTableManager(db *sql.DB, dbType DatabaseType, config TableConfig) *TableManager {
	return &TableManager{
		dbType:      dbType,
		tableConfig: config,
		db:          db,
	}
}

// CreateTable creates the blacklist table if it doesn't exist.
func (tm *TableManager) CreateTable(ctx context.Context) error {
	exists, err := tm.tableExists(ctx)
	if err != nil {
		return errors.Wrap(err, "failed to check if table exists")
	}

	if exists {
		return nil // Table already exists
	}

	createSQL, err := tm.getCreateTableSQL()
	if err != nil {
		return err
	}

	_, err = tm.db.ExecContext(ctx, createSQL)
	if err != nil {
		return errors.Wrapf(err, "failed to create table %s", tm.tableConfig.TableName)
	}

	return nil
}

// tableExists checks if the blacklist table exists in the database.
func (tm *TableManager) tableExists(ctx context.Context) (bool, error) {
	tableName := tm.tableConfig.TableName
	var query string
	var args []interface{}

	switch tm.dbType.Type {
	case DatabasePostgreSQL:
		query = `
			SELECT 1 FROM information_schema.tables 
			WHERE table_schema = 'public' 
			AND table_name = $1
		`
		args = []interface{}{tableName}
	case DatabaseMySQL, DatabaseMariaDB:
		query = `
			SELECT 1 FROM information_schema.tables 
			WHERE table_schema = DATABASE() 
			AND table_name = ?
		`
		args = []interface{}{tableName}
	case DatabaseSQLite:
		query = `
			SELECT 1 FROM sqlite_master 
			WHERE type = 'table' 
			AND name = ?
		`
		args = []interface{}{tableName}
	default:
		return false, errors.Errorf("unsupported database type: %s", tm.dbType.Type)
	}

	rows, err := tm.db.QueryContext(ctx, query, args...)
	if err != nil {
		return false, errors.Wrap(err, "failed to check table existence")
	}
	defer rows.Close()

	return rows.Next(), nil
}

// getCreateTableSQL returns the CREATE TABLE statement for the given database type.
func (tm *TableManager) getCreateTableSQL() (string, error) {
	tableName := tm.tableConfig.TableName

	switch tm.dbType.Type {
	case DatabasePostgreSQL:
		return fmt.Sprintf(`
			CREATE TABLE IF NOT EXISTS %s (
				jti UUID PRIMARY KEY,
				exp BIGINT NOT NULL,
				sub INTEGER NOT NULL,
				created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
			);
			CREATE INDEX IF NOT EXISTS idx_%s_exp ON %s(exp);
			CREATE INDEX IF NOT EXISTS idx_%s_sub ON %s(sub);
		`, tableName, tableName, tableName, tableName, tableName), nil

	case DatabaseMySQL, DatabaseMariaDB:
		return fmt.Sprintf(`
			CREATE TABLE IF NOT EXISTS %s (
				jti CHAR(36) PRIMARY KEY,
				exp BIGINT NOT NULL,
				sub INT NOT NULL,
				created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
				INDEX idx_exp (exp),
				INDEX idx_sub (sub)
			) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
		`, tableName), nil

	case DatabaseSQLite:
		return fmt.Sprintf(`
			CREATE TABLE IF NOT EXISTS %s (
				jti TEXT PRIMARY KEY,
				exp INTEGER NOT NULL,
				sub INTEGER NOT NULL,
				created_at INTEGER DEFAULT (strftime('%%s', 'now'))
			);
			CREATE INDEX IF NOT EXISTS idx_%s_exp ON %s(exp);
			CREATE INDEX IF NOT EXISTS idx_%s_sub ON %s(sub);
		`, tableName, tableName, tableName, tableName, tableName), nil

	default:
		return "", errors.Errorf("unsupported database type: %s", tm.dbType.Type)
	}
}

// SetupAutoCleanup configures database-native automatic cleanup of expired tokens.
func (tm *TableManager) SetupAutoCleanup(ctx context.Context) error {
	if !tm.tableConfig.EnableAutoCleanup {
		return nil
	}

	switch tm.dbType.Type {
	case DatabasePostgreSQL:
		return tm.setupPostgreSQLCleanup(ctx)
	case DatabaseMySQL, DatabaseMariaDB:
		return tm.setupMySQLCleanup(ctx)
	case DatabaseSQLite:
		// SQLite doesn't support automatic cleanup
		return nil
	default:
		return errors.Errorf("unsupported database type: %s", tm.dbType.Type)
	}
}

// setupPostgreSQLCleanup creates a cleanup function for PostgreSQL.
// Note: This creates a function but does not schedule it. You need to use pg_cron
// or an external scheduler to call this function periodically.
func (tm *TableManager) setupPostgreSQLCleanup(ctx context.Context) error {
	tableName := tm.tableConfig.TableName
	functionName := fmt.Sprintf("cleanup_%s", tableName)

	createFunctionSQL := fmt.Sprintf(`
		CREATE OR REPLACE FUNCTION %s()
		RETURNS void AS $$
		BEGIN
			DELETE FROM %s WHERE exp < EXTRACT(EPOCH FROM NOW());
		END;
		$$ LANGUAGE plpgsql;
	`, functionName, tableName)

	_, err := tm.db.ExecContext(ctx, createFunctionSQL)
	if err != nil {
		return errors.Wrap(err, "failed to create cleanup function")
	}

	// Note: Actual scheduling requires pg_cron extension or external tools
	// Users should call this function periodically using:
	// SELECT cleanup_jwtblacklist();
	return nil
}

// setupMySQLCleanup creates a MySQL event for automatic cleanup.
// Note: Requires event_scheduler to be enabled in MySQL/MariaDB configuration.
func (tm *TableManager) setupMySQLCleanup(ctx context.Context) error {
	tableName := tm.tableConfig.TableName
	eventName := fmt.Sprintf("cleanup_%s_event", tableName)
	interval := tm.tableConfig.CleanupInterval

	// Drop existing event if it exists
	dropEventSQL := fmt.Sprintf("DROP EVENT IF EXISTS %s", eventName)
	_, err := tm.db.ExecContext(ctx, dropEventSQL)
	if err != nil {
		return errors.Wrap(err, "failed to drop existing event")
	}

	// Create new event
	createEventSQL := fmt.Sprintf(`
		CREATE EVENT %s
		ON SCHEDULE EVERY %d HOUR
		DO
		DELETE FROM %s WHERE exp < UNIX_TIMESTAMP()
	`, eventName, interval, tableName)

	_, err = tm.db.ExecContext(ctx, createEventSQL)
	if err != nil {
		return errors.Wrapf(err, "failed to create cleanup event (ensure event_scheduler is enabled)")
	}

	return nil
}