EvoBGP/internal/maintenance/policy_executor.go

package maintenance

import (
	"context"
	"fmt"
	"hash/fnv"
	"strings"
	"time"

	"evobgp/internal/observability"
	"evobgp/internal/pgmonitor"
	"evobgp/internal/store"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// PolicyExecutor runs maintenance policies against PostgreSQL.
type PolicyExecutor struct {
	Store store.Backend
	Pool  *pgxpool.Pool
}

// Execute runs cleanup and/or vacuum steps for a policy.
func (e *PolicyExecutor) Execute(ctx context.Context, policy *store.MaintenancePolicy, dryRun bool) (map[string]any, error) {
	start := time.Now()
	if policy == nil {
		return nil, store.ErrInvalidInput
	}
	action := policyAction(policy)
	record := func(status string, detail map[string]any) {
		observability.RecordMaintenancePolicyRun(policy.ID, action, status, dryRun, time.Since(start), rowsDeletedFromDetail(detail))
	}

	if e == nil || e.Pool == nil {
		record("failed", nil)
		return nil, fmt.Errorf("maintenance: postgres not configured")
	}
	if err := ValidateTableName(policy.TableName); err != nil {
		record("failed", nil)
		return nil, err
	}
	if err := ValidateCondition(policy.Condition); err != nil {
		record("failed", nil)
		return nil, err
	}
	if !store.ValidVacuumStrategy(policy.VacuumStrategy) {
		record("failed", nil)
		return nil, store.ErrInvalidInput
	}

	detail := map[string]any{
		"policy_id": policy.ID,
		"table":     policy.TableName,
		"dry_run":   dryRun,
	}

	if policy.RetentionPeriodSec != nil || policy.MaxRows != nil {
		cleanupDetail, err := e.runCleanup(ctx, policy, dryRun)
		for k, v := range cleanupDetail {
			detail[k] = v
		}
		if err != nil {
			record("failed", detail)
			return detail, err
		}
	}

	if policy.VacuumStrategy != store.VacuumStrategyNone {
		kind := vacuumKind(policy.VacuumStrategy)
		vacDetail, err := pgmonitor.ExecMaintenance(ctx, e.Pool, kind, policy.TableName, dryRun)
		if vacDetail != nil {
			detail["vacuum"] = vacDetail
		}
		if err != nil {
			record("failed", detail)
			return detail, err
		}
	}

	if !dryRun && e.Store != nil {
		_ = e.Store.TouchMaintenancePolicyRun(policy.ID, "succeeded", "")
	}
	record("succeeded", detail)
	return detail, nil
}

func policyAction(p *store.MaintenancePolicy) string {
	if p == nil {
		return "run"
	}
	if p.RetentionPeriodSec != nil || p.MaxRows != nil {
		if p.VacuumStrategy != store.VacuumStrategyNone {
			return "cleanup_vacuum"
		}
		return "cleanup"
	}
	if p.VacuumStrategy != store.VacuumStrategyNone {
		return p.VacuumStrategy
	}
	return "run"
}

func rowsDeletedFromDetail(detail map[string]any) int64 {
	if detail == nil {
		return 0
	}
	switch v := detail["deleted"].(type) {
	case int64:
		return v
	case int:
		return int64(v)
	case float64:
		return int64(v)
	default:
		return 0
	}
}

func (e *PolicyExecutor) runCleanup(ctx context.Context, policy *store.MaintenancePolicy, dryRun bool) (map[string]any, error) {
	detail := map[string]any{"cleanup": true}
	limit := NormalizeBatchLimit(policy.MaxRows)
	qualTable := pgx.Identifier{policy.TableName}.Sanitize()
	cond := store.NormalizeMaintenancePolicyCondition(policy.Condition)

	tx, err := e.Pool.Begin(ctx)
	if err != nil {
		return detail, fmt.Errorf("maintenance: begin tx: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	lockKey := advisoryKey(policy.ID)
	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, lockKey); err != nil {
		return detail, fmt.Errorf("maintenance: advisory lock: %w", err)
	}
	if _, err := tx.Exec(ctx, fmt.Sprintf(`SET LOCAL statement_timeout = '%ds'`, DefaultStatementTimeoutSec)); err != nil {
		return detail, fmt.Errorf("maintenance: statement_timeout: %w", err)
	}

	var args []any
	where := cond
	argN := 1
	if policy.RetentionPeriodSec != nil && *policy.RetentionPeriodSec > 0 {
		cutoff := time.Now().UTC().Add(-time.Duration(*policy.RetentionPeriodSec) * time.Second)
		where = fmt.Sprintf("(%s) AND created_at < $%d", cond, argN)
		args = append(args, cutoff)
		argN++
	}

	countSQL := fmt.Sprintf(`SELECT count(*) FROM %s WHERE %s`, qualTable, where)
	var wouldDelete int64
	if err := tx.QueryRow(ctx, countSQL, args...).Scan(&wouldDelete); err != nil {
		return detail, fmt.Errorf("maintenance: count: %w", err)
	}
	detail["would_delete"] = wouldDelete
	if dryRun {
		return detail, nil
	}

	deleteSQL := fmt.Sprintf(`
		DELETE FROM %s WHERE ctid IN (
			SELECT ctid FROM %s WHERE %s LIMIT $%d
		)`, qualTable, qualTable, where, argN)
	args = append(args, limit)
	tag, err := tx.Exec(ctx, deleteSQL, args...)
	if err != nil {
		return detail, fmt.Errorf("maintenance: delete: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return detail, fmt.Errorf("maintenance: commit: %w", err)
	}
	detail["deleted"] = tag.RowsAffected()
	return detail, nil
}

func vacuumKind(strategy string) string {
	switch strings.TrimSpace(strategy) {
	case store.VacuumStrategyVacuum:
		return "vacuum"
	case store.VacuumStrategyAnalyze:
		return "analyze"
	case store.VacuumStrategyVacuumAnalyze:
		return "vacuum_analyze"
	case store.VacuumStrategyReindex:
		return "reindex"
	default:
		return "vacuum"
	}
}

func advisoryKey(policyID string) int64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte("maint:" + policyID))
	return int64(h.Sum64())
}