PriceForge/internal/dbutil/timeout.go

package dbutil

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"gorm.io/gorm"
)

// QueryWithTimeout executes a database query with timeout and optional retry
type QueryWithTimeout struct {
	DB            *gorm.DB
	Timeout       time.Duration
	RetryAttempts int
	RetryDelay    time.Duration
}

// DefaultQueryTimeout creates a query executor with sensible defaults
func DefaultQueryTimeout(db *gorm.DB) *QueryWithTimeout {
	return &QueryWithTimeout{
		DB:            db,
		Timeout:       10 * time.Second,
		RetryAttempts: 1, // Total 2 attempts (initial + 1 retry)
		RetryDelay:    100 * time.Millisecond,
	}
}

// WithTimeout creates a query executor with custom timeout
func WithTimeout(db *gorm.DB, timeout time.Duration) *QueryWithTimeout {
	return &QueryWithTimeout{
		DB:            db,
		Timeout:       timeout,
		RetryAttempts: 1,
		RetryDelay:    100 * time.Millisecond,
	}
}

// Execute runs the query with timeout and retry logic
func (q *QueryWithTimeout) Execute(queryFn func(*gorm.DB) error) error {
	var lastErr error

	for attempt := 0; attempt <= q.RetryAttempts; attempt++ {
		if attempt > 0 {
			// Wait before retry
			slog.Warn("DB query retry", "attempt", attempt, "timeout", q.Timeout)
			fmt.Printf("[DB] ⚠️  Retry attempt %d/%d (timeout: %v)\n", attempt, q.RetryAttempts, q.Timeout)
			time.Sleep(q.RetryDelay)
		}

		ctx, cancel := context.WithTimeout(context.Background(), q.Timeout)
		dbWithTimeout := q.DB.WithContext(ctx)

		startTime := time.Now()
		err := queryFn(dbWithTimeout)
		duration := time.Since(startTime)
		cancel()

		if err == nil {
			if duration > q.Timeout/2 {
				// Query took more than half the timeout - warning
				slog.Warn("Slow DB query", "duration", duration, "timeout", q.Timeout)
				fmt.Printf("[DB] ⚠️  Slow query: %v (timeout: %v)\n", duration, q.Timeout)
			}
			return nil
		}

		// Log the error
		if ctx.Err() == context.DeadlineExceeded {
			slog.Error("DB query timeout", "timeout", q.Timeout, "attempt", attempt, "error", err)
			fmt.Printf("[DB] ❌ Query TIMEOUT after %v (attempt %d/%d): %v\n", q.Timeout, attempt+1, q.RetryAttempts+1, err)
		} else {
			slog.Error("DB query error", "duration", duration, "attempt", attempt, "error", err)
			fmt.Printf("[DB] ❌ Query ERROR after %v (attempt %d/%d): %v\n", duration, attempt+1, q.RetryAttempts+1, err)
		}

		lastErr = err

		// Don't retry on certain errors (e.g., record not found)
		if err == gorm.ErrRecordNotFound {
			return err
		}
	}

	fmt.Printf("[DB] 🔴 Query FAILED after all retries: %v\n", lastErr)
	return lastErr
}

// ExecuteRaw is a convenience method for Raw queries
func (q *QueryWithTimeout) ExecuteRaw(query string, args ...interface{}) *gorm.DB {
	ctx, cancel := context.WithTimeout(context.Background(), q.Timeout)
	defer cancel()
	return q.DB.WithContext(ctx).Raw(query, args...)
}