feat: add stdin support and retry logic for all search commandsHEADv1.1.1v0.1.14master

1 files changed, 133 insertions, 43 deletions

diff --git a/internal/client/client.go b/internal/client/client.go
index 4aa6f06..4ca61e4 100644
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -4,11 +4,16 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"git.db.org.ai/dborg/internal/config"
 	"io"
+	"math/rand"
 	"net/http"
 	"net/url"
+	"os"
+	"strconv"
+	"strings"
 	"time"
+
+	"git.db.org.ai/dborg/internal/config"
 )
 
 type Client struct {
@@ -20,93 +25,178 @@ func New(cfg *config.Config) (*Client, error) {
 	if err := cfg.Validate(); err != nil {
 		return nil, err
 	}
-
-	return &Client{
-		config: cfg,
-		httpClient: &http.Client{
-			Timeout: cfg.Timeout,
-		},
-	}, nil
+	return newClient(cfg), nil
 }
 
 func NewUnauthenticated(cfg *config.Config) (*Client, error) {
+	return newClient(cfg), nil
+}
+
+func newClient(cfg *config.Config) *Client {
 	return &Client{
 		config: cfg,
 		httpClient: &http.Client{
 			Timeout: cfg.Timeout,
 		},
-	}, nil
+	}
+}
+
+func (c *Client) debugf(format string, args ...interface{}) {
+	if !c.config.Debug {
+		return
+	}
+	fmt.Fprintf(os.Stderr, "[dborg] "+format+"\n", args...)
+}
+
+func redactKey(key string) string {
+	if len(key) <= 8 {
+		return "***"
+	}
+	return key[:4] + "..." + key[len(key)-4:]
+}
+
+func isRetryable(statusCode int) bool {
+	return statusCode == http.StatusTooManyRequests ||
+		statusCode == http.StatusRequestTimeout ||
+		statusCode >= 500
+}
+
+func backoffDelay(attempt int, retryAfter string) time.Duration {
+	if retryAfter != "" {
+		if secs, err := strconv.Atoi(strings.TrimSpace(retryAfter)); err == nil && secs > 0 {
+			return time.Duration(secs) * time.Second
+		}
+		if t, err := http.ParseTime(retryAfter); err == nil {
+			if d := time.Until(t); d > 0 {
+				return d
+			}
+		}
+	}
+	base := time.Duration(1<<attempt) * time.Second
+	if base > 30*time.Second {
+		base = 30 * time.Second
+	}
+	jitter := time.Duration(rand.Int63n(int64(base) / 2))
+	return base + jitter
 }
 
 func (c *Client) doRequest(method, path string, params url.Values, body interface{}) ([]byte, error) {
 	fullURL := c.config.BaseURL + path
-	if params != nil && len(params) > 0 {
+	if len(params) > 0 {
 		fullURL += "?" + params.Encode()
 	}
 
-	var reqBody io.Reader
+	var bodyBytes []byte
 	if body != nil {
-		jsonData, err := json.Marshal(body)
+		var err error
+		bodyBytes, err = json.Marshal(body)
 		if err != nil {
 			return nil, fmt.Errorf("failed to marshal request body: %w", err)
 		}
-		reqBody = bytes.NewBuffer(jsonData)
-	}
-
-	req, err := http.NewRequest(method, fullURL, reqBody)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
-	req.Header.Set("X-API-Key", c.config.APIKey)
-	req.Header.Set("User-Agent", c.config.UserAgent)
-	if body != nil {
-		req.Header.Set("Content-Type", "application/json")
+	c.debugf("→ %s %s (api_key=%s)", method, fullURL, redactKey(c.config.APIKey))
+	if len(bodyBytes) > 0 {
+		c.debugf("  body: %s", string(bodyBytes))
 	}
 
-	var resp *http.Response
 	var lastErr error
 
 	for attempt := 0; attempt <= c.config.MaxRetries; attempt++ {
 		if attempt > 0 {
-			time.Sleep(time.Duration(attempt) * time.Second)
+			delay := backoffDelay(attempt, "")
+			if lastErr != nil {
+				if ra, ok := retryAfterFromErr(lastErr); ok {
+					delay = backoffDelay(attempt, ra)
+				}
+			}
+			c.debugf("  retry %d/%d after %s (last error: %v)", attempt, c.config.MaxRetries, delay, lastErr)
+			time.Sleep(delay)
 		}
 
-		resp, err = c.httpClient.Do(req)
+		var reqBody io.Reader
+		if bodyBytes != nil {
+			reqBody = bytes.NewReader(bodyBytes)
+		}
+
+		req, err := http.NewRequest(method, fullURL, reqBody)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create request: %w", err)
+		}
+
+		req.Header.Set("X-API-Key", c.config.APIKey)
+		req.Header.Set("User-Agent", c.config.UserAgent)
+		if bodyBytes != nil {
+			req.Header.Set("Content-Type", "application/json")
+		}
+
+		start := time.Now()
+		resp, err := c.httpClient.Do(req)
 		if err != nil {
 			lastErr = err
+			c.debugf("← network error after %s: %v", time.Since(start), err)
 			continue
 		}
 
-		defer resp.Body.Close()
+		respBody, readErr := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		c.debugf("← %d %s (%s, %d bytes)", resp.StatusCode, http.StatusText(resp.StatusCode), time.Since(start), len(respBody))
 
-		if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
-			return io.ReadAll(resp.Body)
+		if readErr != nil {
+			lastErr = fmt.Errorf("failed to read response body: %w", readErr)
+			if !isRetryable(resp.StatusCode) {
+				return nil, lastErr
+			}
+			continue
 		}
 
-		bodyBytes, _ := io.ReadAll(resp.Body)
-
-		switch resp.StatusCode {
-		case http.StatusForbidden:
-			lastErr = fmt.Errorf("access denied (403): %s - This endpoint requires premium access", string(bodyBytes))
-		case http.StatusUnauthorized:
-			lastErr = fmt.Errorf("unauthorized (401): %s - Check your API key", string(bodyBytes))
-		case http.StatusTooManyRequests:
-			lastErr = fmt.Errorf("rate limit exceeded (429): %s", string(bodyBytes))
-		case http.StatusBadRequest:
-			lastErr = fmt.Errorf("bad request (400): %s", string(bodyBytes))
-		default:
-			lastErr = fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
+		if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
+			return respBody, nil
 		}
 
-		if resp.StatusCode != http.StatusTooManyRequests && resp.StatusCode < 500 {
-			break
+		lastErr = httpError(resp.StatusCode, respBody, resp.Header.Get("Retry-After"))
+
+		if !isRetryable(resp.StatusCode) {
+			return nil, lastErr
 		}
 	}
 
 	return nil, lastErr
 }
 
+type apiError struct {
+	status     int
+	message    string
+	retryAfter string
+}
+
+func (e *apiError) Error() string { return e.message }
+
+func retryAfterFromErr(err error) (string, bool) {
+	if ae, ok := err.(*apiError); ok && ae.retryAfter != "" {
+		return ae.retryAfter, true
+	}
+	return "", false
+}
+
+func httpError(status int, body []byte, retryAfter string) error {
+	msg := string(body)
+	var formatted string
+	switch status {
+	case http.StatusForbidden:
+		formatted = fmt.Sprintf("access denied (403): %s - This endpoint requires premium access", msg)
+	case http.StatusUnauthorized:
+		formatted = fmt.Sprintf("unauthorized (401): %s - Check your API key", msg)
+	case http.StatusTooManyRequests:
+		formatted = fmt.Sprintf("rate limit exceeded (429): %s", msg)
+	case http.StatusBadRequest:
+		formatted = fmt.Sprintf("bad request (400): %s", msg)
+	default:
+		formatted = fmt.Sprintf("API request failed with status %d: %s", status, msg)
+	}
+	return &apiError{status: status, message: formatted, retryAfter: retryAfter}
+}
+
 func (c *Client) Get(path string, params url.Values) ([]byte, error) {
 	return c.doRequest(http.MethodGet, path, params, nil)
 }