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...) }