Add query timeout to prevent indefinite hangs on S3 disconnection

cmd/arc/main.go

@@ -728,8 +728,8 @@ func main() {
 	}
 	lineProtocolHandler.RegisterRoutes(server.GetApp())
 
-	// Register Query handler
-	queryHandler := api.NewQueryHandler(db, storageBackend, logger.Get("query"))
+	// Register Query handler with dedicated query timeout
+	queryHandler := api.NewQueryHandler(db, storageBackend, logger.Get("query"), cfg.Query.Timeout)
 	if authManager != nil && rbacManager != nil {
 		queryHandler.SetAuthAndRBAC(authManager, rbacManager)
 	}

internal/api/query.go

@@ -456,6 +456,7 @@ type QueryHandler struct {
 	rbacManager      RBACChecker
 	debugEnabled     bool // Cached check for debug logging to avoid repeated level checks
 	parallelExecutor *query.ParallelExecutor
+	queryTimeout     time.Duration // Query timeout (0 = no timeout)
 
 	// Cluster routing support
 	router *cluster.Router
@@ -587,10 +588,18 @@ func validateWhereClauseQuery(where string) error {
 }
 
 // NewQueryHandler creates a new query handler
-func NewQueryHandler(db *database.DuckDB, storage storage.Backend, logger zerolog.Logger) *QueryHandler {
+// queryTimeoutSeconds: timeout for query execution in seconds (0 = no timeout)
+func NewQueryHandler(db *database.DuckDB, storage storage.Backend, logger zerolog.Logger, queryTimeoutSeconds int) *QueryHandler {
 	handlerLogger := logger.With().Str("component", "query-handler").Logger()
 	pruner := pruning.NewPartitionPruner(logger)
 	pruner.SetStorageBackend(storage) // Enable S3/Azure partition filtering
+
+	var queryTimeout time.Duration
+	if queryTimeoutSeconds > 0 {
+		queryTimeout = time.Duration(queryTimeoutSeconds) * time.Second
+		handlerLogger.Info().Int("timeout_seconds", queryTimeoutSeconds).Msg("Query timeout configured")
+	}
+
 	return &QueryHandler{
 		db:               db,
 		storage:          storage,
@@ -601,6 +610,7 @@ func NewQueryHandler(db *database.DuckDB, storage storage.Backend, logger zerolo
 		rbacManager:      nil,
 		debugEnabled:     handlerLogger.GetLevel() <= zerolog.DebugLevel,
 		parallelExecutor: query.NewParallelExecutor(db.DB(), query.DefaultParallelConfig(), handlerLogger),
+		queryTimeout:     queryTimeout,
 	}
 }
 
@@ -1050,15 +1060,34 @@ localProcessing:
 		var rows *sql.Rows
 		var err error
 
+		// Create context with timeout if configured (0 = no timeout)
+		ctx := c.UserContext()
+		var cancel context.CancelFunc
+		if h.queryTimeout > 0 {
+			ctx, cancel = context.WithTimeout(ctx, h.queryTimeout)
+			defer cancel()
+		}
+
 		if profileMode {
-			// Use profiled query to capture timing breakdown
-			rows, profile, err = h.db.QueryWithProfile(convertedSQL)
+			// Use profiled query to capture timing breakdown (with timeout support)
+			rows, profile, err = h.db.QueryWithProfileContext(ctx, convertedSQL)
 		} else {
-			rows, err = h.db.Query(convertedSQL)
+			rows, err = h.db.QueryContext(ctx, convertedSQL)
 		}
 
 		if err != nil {
 			m.IncQueryErrors()
+			// Check if it was a timeout
+			if h.queryTimeout > 0 && ctx.Err() == context.DeadlineExceeded {
+				m.IncQueryTimeouts()
+				h.logger.Error().Err(err).Str("sql", req.SQL).Dur("timeout", h.queryTimeout).Msg("Query timed out")
+				return c.Status(fiber.StatusGatewayTimeout).JSON(QueryResponse{
+					Success:         false,
+					Error:           "Query timed out",
+					ExecutionTimeMs: float64(time.Since(start).Milliseconds()),
+					Timestamp:       timestamp,
+				})
+			}
 			h.logger.Error().Err(err).Str("sql", req.SQL).Msg("Query execution failed")
 			return c.Status(fiber.StatusInternalServerError).JSON(QueryResponse{
 				Success:         false,
@@ -2250,9 +2279,30 @@ func (h *QueryHandler) estimateQuery(c *fiber.Ctx) error {
 		Str("count_sql", countSQL).
 		Msg("Estimating query")
 
-	// Execute count query
-	rows, err := h.db.Query(countSQL)
+	m := metrics.Get()
+
+	// Create context with timeout if configured
+	ctx := c.UserContext()
+	var cancel context.CancelFunc
+	if h.queryTimeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, h.queryTimeout)
+		defer cancel()
+	}
+
+	// Execute count query with timeout support
+	rows, err := h.db.QueryContext(ctx, countSQL)
 	if err != nil {
+		// Check if it was a timeout
+		if h.queryTimeout > 0 && ctx.Err() == context.DeadlineExceeded {
+			m.IncQueryTimeouts()
+			h.logger.Error().Err(err).Str("sql", countSQL).Dur("timeout", h.queryTimeout).Msg("Estimate query timed out")
+			return c.Status(fiber.StatusGatewayTimeout).JSON(EstimateResponse{
+				Success:         false,
+				Error:           "Query timed out",
+				WarningLevel:    "error",
+				ExecutionTimeMs: float64(time.Since(start).Milliseconds()),
+			})
+		}
 		h.logger.Error().Err(err).Str("sql", countSQL).Msg("Estimate query failed")
 		return c.JSON(EstimateResponse{
 			Success:         false,

internal/api/query_arrow.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -75,9 +76,26 @@ func (h *QueryHandler) executeQueryArrow(c *fiber.Ctx) error {
 		Str("header_db", headerDB).
 		Msg("Executing Arrow query")
 
-	// Execute query using standard database/sql interface
-	rows, err := h.db.Query(convertedSQL)
+	// Create context with timeout if configured
+	ctx := c.UserContext()
+	var cancel context.CancelFunc
+	if h.queryTimeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, h.queryTimeout)
+		defer cancel()
+	}
+
+	// Execute query using standard database/sql interface with timeout support
+	rows, err := h.db.QueryContext(ctx, convertedSQL)
 	if err != nil {
+		// Check if it was a timeout
+		if h.queryTimeout > 0 && ctx.Err() == context.DeadlineExceeded {
+			m.IncQueryTimeouts()
+			h.logger.Error().Err(err).Str("sql", req.SQL).Dur("timeout", h.queryTimeout).Msg("Arrow query timed out")
+			return c.Status(fiber.StatusGatewayTimeout).JSON(fiber.Map{
+				"success": false,
+				"error":   "Query timed out",
+			})
+		}
 		m.IncQueryErrors()
 		h.logger.Error().Err(err).Str("sql", req.SQL).Msg("Arrow query execution failed")
 		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{

internal/config/config.go

@@ -165,6 +165,7 @@ type MQTTConfig struct {
 
 // QueryConfig holds configuration for query execution optimizations
 type QueryConfig struct {
+	Timeout           int   // Query execution timeout in seconds (0 = no timeout, default: 300)
 	EnableS3Cache     bool  // Enable S3 file caching for faster repeated reads (useful for CTEs/subqueries)
 	S3CacheSize       int64 // Cache size in bytes (parsed from "128MB", "256MB", etc.)
 	S3CacheTTLSeconds int   // Cache entry TTL in seconds (default: 3600 = 1 hour)
@@ -440,6 +441,7 @@ func Load() (*Config, error) {
 			Enabled: v.GetBool("mqtt.enabled"),
 		},
 		Query: QueryConfig{
+			Timeout:           v.GetInt("query.timeout"),
 			EnableS3Cache:     v.GetBool("query.enable_s3_cache"),
 			S3CacheSize:       s3CacheSize,
 			S3CacheTTLSeconds: v.GetInt("query.s3_cache_ttl_seconds"),
@@ -623,6 +625,7 @@ func setDefaults(v *viper.Viper) {
 	v.SetDefault("mqtt.enabled", false) // Feature toggle only - disabled by default
 
 	// Query defaults
+	v.SetDefault("query.timeout", 300)                 // 5 minute query timeout (0 = no timeout)
 	v.SetDefault("query.enable_s3_cache", false)       // Disabled by default (opt-in feature)
 	v.SetDefault("query.s3_cache_size", "128MB")       // 128MB cache (256 blocks × 512KB)
 	v.SetDefault("query.s3_cache_ttl_seconds", 3600)   // 1 hour

internal/database/duckdb.go

@@ -1,6 +1,7 @@
 package database
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"fmt"
@@ -398,6 +399,29 @@ func (d *DuckDB) Query(query string, args ...interface{}) (*sql.Rows, error) {
 	return rows, nil
 }
 
+// QueryContext executes a query with context support for timeout/cancellation
+func (d *DuckDB) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
+	start := time.Now()
+	rows, err := d.db.QueryContext(ctx, query, args...)
+	elapsed := time.Since(start)
+
+	if err != nil {
+		d.logger.Error().
+			Err(err).
+			Str("query", query).
+			Dur("elapsed", elapsed).
+			Msg("Query failed")
+		return nil, fmt.Errorf("query failed: %w", err)
+	}
+
+	d.logger.Debug().
+		Str("query", query).
+		Dur("elapsed", elapsed).
+		Msg("Query executed")
+
+	return rows, nil
+}
+
 // Exec executes a statement without returning rows
 func (d *DuckDB) Exec(query string, args ...interface{}) (sql.Result, error) {
 	start := time.Now()
@@ -445,11 +469,17 @@ func (d *DuckDB) DB() *sql.DB {
 // QueryWithProfile executes a query and returns timing breakdown using DuckDB profiling
 // This is used to measure parsing/planning overhead for optimization decisions
 func (d *DuckDB) QueryWithProfile(query string) (*sql.Rows, *QueryProfile, error) {
+	return d.QueryWithProfileContext(context.Background(), query)
+}
+
+// QueryWithProfileContext executes a query with context support for timeout/cancellation
+// and returns timing breakdown using DuckDB profiling
+func (d *DuckDB) QueryWithProfileContext(ctx context.Context, query string) (*sql.Rows, *QueryProfile, error) {
 	// Create a temporary file for profiling output
 	tmpFile, err := os.CreateTemp("", "duckdb_profile_*.json")
 	if err != nil {
 		// Fall back to regular query if we can't create temp file
-		rows, err := d.Query(query)
+		rows, err := d.QueryContext(ctx, query)
 		return rows, nil, err
 	}
 	profilePath := tmpFile.Name()
@@ -468,9 +498,9 @@ func (d *DuckDB) QueryWithProfile(query string) (*sql.Rows, *QueryProfile, error
 		d.logger.Warn().Err(err).Msg("Failed to set custom profiling settings")
 	}
 
-	// Execute the query with timing
+	// Execute the query with timing and context
 	start := time.Now()
-	rows, err := d.db.Query(query)
+	rows, err := d.db.QueryContext(ctx, query)
 	totalTime := time.Since(start)
 
 	// Disable profiling

internal/metrics/metrics.go

@@ -44,6 +44,7 @@ type Metrics struct {
 	queryRequestsTotal    atomic.Int64
 	querySuccessTotal     atomic.Int64
 	queryErrorsTotal      atomic.Int64
+	queryTimeoutsTotal    atomic.Int64
 	queryRowsTotal        atomic.Int64
 	queryLatencySum       atomic.Int64 // microseconds
 	queryLatencyCount     atomic.Int64
@@ -179,6 +180,7 @@ func (m *Metrics) IncLineProtocolBytes(bytes int64)   { m.lineprotocolBytesTotal
 func (m *Metrics) IncQueryRequests()             { m.queryRequestsTotal.Add(1) }
 func (m *Metrics) IncQuerySuccess()              { m.querySuccessTotal.Add(1) }
 func (m *Metrics) IncQueryErrors()               { m.queryErrorsTotal.Add(1) }
+func (m *Metrics) IncQueryTimeouts()             { m.queryTimeoutsTotal.Add(1) }
 func (m *Metrics) IncQueryRows(count int64)      { m.queryRowsTotal.Add(count) }
 
 // RecordQueryLatency records query latency in microseconds
@@ -288,6 +290,7 @@ func (m *Metrics) Snapshot() map[string]interface{} {
 		"query_requests_total":  m.queryRequestsTotal.Load(),
 		"query_success_total":   m.querySuccessTotal.Load(),
 		"query_errors_total":    m.queryErrorsTotal.Load(),
+		"query_timeouts_total":  m.queryTimeoutsTotal.Load(),
 		"query_rows_total":      m.queryRowsTotal.Load(),
 		"query_latency_sum_us":  m.queryLatencySum.Load(),
 		"query_latency_count":   m.queryLatencyCount.Load(),
@@ -445,6 +448,10 @@ func (m *Metrics) PrometheusFormat() string {
 	b = append(b, "# TYPE arc_query_errors_total counter\n"...)
 	b = appendMetric(b, "arc_query_errors_total", float64(m.queryErrorsTotal.Load()))
 
+	b = append(b, "# HELP arc_query_timeouts_total Queries that exceeded timeout\n"...)
+	b = append(b, "# TYPE arc_query_timeouts_total counter\n"...)
+	b = appendMetric(b, "arc_query_timeouts_total", float64(m.queryTimeoutsTotal.Load()))
+
 	b = append(b, "# HELP arc_query_rows_total Total rows returned by queries\n"...)
 	b = append(b, "# TYPE arc_query_rows_total counter\n"...)
 	b = appendMetric(b, "arc_query_rows_total", float64(m.queryRowsTotal.Load()))