fix(clickhouse): enforce time and read limits on every query

cmd/tools/generate-implementation/endpoints.go

@@ -93,7 +93,7 @@ func (s *Server) %s(w http.ResponseWriter, r *http.Request, params handlers.%s)
 	if err != nil {
 		span.RecordError(err)
 		span.SetStatus(codes.Error, "query failed")
-		writeError(w, http.StatusInternalServerError, err)
+		apierrors.FromQuery(err).WriteJSON(w)
 		return
 	}
 	defer rows.Close()
@@ -113,6 +113,15 @@ func (s *Server) %s(w http.ResponseWriter, r *http.Request, params handlers.%s)
 		}
 		items = append(items, item)
 	}
+	// A limit can stop the query after the first rows arrived; return the error, not a short page.
+	if err := rows.Err(); err != nil {
+		scanSpan.RecordError(err)
+		scanSpan.End()
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "query failed")
+		apierrors.FromQuery(err).WriteJSON(w)
+		return
+	}
 	scanSpan.SetAttributes(attribute.Int("result.count", len(items)))
 	scanSpan.End()
 
@@ -203,7 +212,7 @@ func (s *Server) %s(w http.ResponseWriter, r *http.Request, %s %s) {
 	if err != nil {
 		span.RecordError(err)
 		span.SetStatus(codes.Error, "query failed")
-		writeError(w, http.StatusInternalServerError, err)
+		apierrors.FromQuery(err).WriteJSON(w)
 		return
 	}
 	defer rows.Close()
@@ -220,6 +229,13 @@ func (s *Server) %s(w http.ResponseWriter, r *http.Request, %s %s) {
 			writeError(w, http.StatusInternalServerError, err)
 			return
 		}
+	} else if err := rows.Err(); err != nil {
+		scanSpan.RecordError(err)
+		scanSpan.End()
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "query failed")
+		apierrors.FromQuery(err).WriteJSON(w)
+		return
 	} else {
 		scanSpan.SetAttributes(attribute.Bool("result.found", false))
 		scanSpan.End()

config.example.yaml

@@ -17,7 +17,9 @@ clickhouse:
   dial_timeout: 10s
   read_timeout: 30s
   write_timeout: 30s
-  max_execution_time: 60
+  max_execution_time: 20
+  max_rows_to_read: 1000000000
+  max_bytes_to_read: 68719476736 # 64 GiB
   insecure_skip_verify: false
   discovery:
     prefixes:

internal/config/config.go

@@ -63,6 +63,11 @@ type ClickHouseConfig struct {
 	WriteTimeout     time.Duration `mapstructure:"write_timeout"`
 	MaxExecutionTime int           `mapstructure:"max_execution_time"`
 
+	// Read limits: a query that would read more rows or bytes fails, and the API returns a 400.
+	// 0 disables a limit.
+	MaxRowsToRead  uint64 `mapstructure:"max_rows_to_read"`
+	MaxBytesToRead uint64 `mapstructure:"max_bytes_to_read"`
+
 	// TLS settings
 	InsecureSkipVerify bool `mapstructure:"insecure_skip_verify"`
 }
@@ -130,7 +135,9 @@ func Load(configFile string) (*Config, error) {
 	viper.SetDefault("clickhouse.dial_timeout", 10*time.Second)
 	viper.SetDefault("clickhouse.read_timeout", 30*time.Second)
 	viper.SetDefault("clickhouse.write_timeout", 30*time.Second)
-	viper.SetDefault("clickhouse.max_execution_time", 60)
+	viper.SetDefault("clickhouse.max_execution_time", 20)
+	viper.SetDefault("clickhouse.max_rows_to_read", 1_000_000_000)
+	viper.SetDefault("clickhouse.max_bytes_to_read", 64<<30)
 	viper.SetDefault("clickhouse.use_final", false)
 	viper.SetDefault("clickhouse.insecure_skip_verify", false)
 

internal/database/clickhouse.go

@@ -38,6 +38,9 @@ type Client struct {
 	conn   driver.Conn
 	log    logrus.FieldLogger
 	config *config.ClickHouseConfig
+	// settings is a SETTINGS clause appended to every query. It travels in the query text, so it reaches
+	// ClickHouse through proxies such as chproxy that drop settings sent as URL parameters.
+	settings string
 }
 
 // Ensure Client implements DatabaseClient interface.
@@ -85,25 +88,58 @@ func NewClient(cfg *config.ClickHouseConfig, logger logrus.FieldLogger) (*Client
 	}).Info("ClickHouse client initialised")
 
 	return &Client{
-		conn:   conn,
-		log:    log,
-		config: cfg,
+		conn:     conn,
+		log:      log,
+		config:   cfg,
+		settings: settingsClause(cfg),
 	}, nil
 }
 
 // Query executes a SQL query and returns rows (native driver interface).
 func (c *Client) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) {
-	return c.conn.Query(ctx, query, args...)
+	return c.conn.Query(ctx, c.withSettings(query), args...)
 }
 
 // QueryRow executes a query that is expected to return at most one row.
 func (c *Client) QueryRow(ctx context.Context, query string, args ...any) driver.Row {
-	return c.conn.QueryRow(ctx, query, args...)
+	return c.conn.QueryRow(ctx, c.withSettings(query), args...)
 }
 
 // Select executes a query and scans results directly into a slice of structs.
 func (c *Client) Select(ctx context.Context, dest any, query string, args ...any) error {
-	return c.conn.Select(ctx, dest, query, args...)
+	return c.conn.Select(ctx, dest, c.withSettings(query), args...)
+}
+
+// settingsClause builds the SETTINGS clause that limits each query's run time and how much it reads.
+func settingsClause(cfg *config.ClickHouseConfig) string {
+	settings := make([]string, 0, 4)
+
+	if cfg.MaxExecutionTime > 0 {
+		settings = append(settings, fmt.Sprintf("max_execution_time = %d", cfg.MaxExecutionTime))
+	}
+
+	if cfg.MaxRowsToRead > 0 {
+		settings = append(settings, fmt.Sprintf("max_rows_to_read = %d", cfg.MaxRowsToRead))
+	}
+
+	if cfg.MaxBytesToRead > 0 {
+		settings = append(settings, fmt.Sprintf("max_bytes_to_read = %d", cfg.MaxBytesToRead))
+	}
+
+	if len(settings) == 0 {
+		return ""
+	}
+
+	return " SETTINGS " + strings.Join(append(settings, "read_overflow_mode = 'throw'"), ", ")
+}
+
+// withSettings appends the client's SETTINGS clause to a query that has none of its own.
+func (c *Client) withSettings(query string) string {
+	if c.settings == "" || strings.Contains(strings.ToUpper(query), " SETTINGS ") {
+		return query
+	}
+
+	return strings.TrimRight(strings.TrimSpace(query), ";") + c.settings
 }
 
 // Exec executes a query without returning any rows.

internal/database/clickhouse_test.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/ethpandaops/cbt-api/internal/config"
+	apierrors "github.com/ethpandaops/cbt-api/internal/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -334,10 +336,12 @@ func TestClient_QueryTimeout(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 	defer cancel()
 
-	var result []uint64
+	var result []struct {
+		V uint8 `ch:"v"`
+	}
 
-	// Query that will take a long time (summing a billion numbers)
-	err := client.Select(ctx, &result, "SELECT sum(number) FROM numbers(1000000000)")
+	// A query that always outlasts the deadline
+	err := client.Select(ctx, &result, "SELECT sleep(2) AS v")
 	require.Error(t, err)
 
 	// Accept either "deadline exceeded" or "i/o timeout" as both indicate timeout
@@ -562,3 +566,51 @@ func createTestClient(t *testing.T) *Client {
 
 	return client
 }
+
+// TestClient_ReadLimits checks that the read limits reach ClickHouse over both protocols: a query that reads more
+// rows than allowed fails with a limit error that the API maps to a 400, and a smaller one succeeds.
+func TestClient_ReadLimits(t *testing.T) {
+	userInfo := fmt.Sprintf("%s:%s", sharedContainer.User, sharedContainer.Password)
+
+	for name, dsn := range map[string]string{
+		"native": fmt.Sprintf("clickhouse://%s@%s", userInfo, net.JoinHostPort(sharedContainerHost, sharedNativePort)),
+		"http":   fmt.Sprintf("http://%s@%s", userInfo, net.JoinHostPort(sharedContainerHost, sharedHTTPPort)),
+	} {
+		t.Run(name, func(t *testing.T) {
+			cfg := &config.ClickHouseConfig{
+				DSN:              dsn,
+				Database:         "default",
+				MaxExecutionTime: 20,
+				MaxRowsToRead:    1000,
+				DialTimeout:      5 * time.Second,
+				ReadTimeout:      10 * time.Second,
+			}
+
+			log := logrus.New()
+			log.SetLevel(logrus.ErrorLevel)
+
+			client, err := NewClient(cfg, log)
+			require.NoError(t, err)
+
+			defer client.Close()
+
+			ctx := context.Background()
+
+			var small []struct {
+				N uint64 `ch:"n"`
+			}
+			require.NoError(t, client.Select(ctx, &small, "SELECT count() AS n FROM numbers(500)"))
+			assert.Equal(t, uint64(500), small[0].N)
+
+			var big []struct {
+				N uint64 `ch:"n"`
+			}
+
+			err = client.Select(ctx, &big, "SELECT count() AS n FROM numbers(100000)")
+			require.Error(t, err)
+
+			status := apierrors.FromQuery(err)
+			assert.Equal(t, http.StatusBadRequest, apierrors.HTTPStatus(status.Code), err.Error())
+		})
+	}
+}

internal/errors/status.go

@@ -2,9 +2,13 @@ package errors
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
+	"regexp"
+	"strconv"
 
+	"github.com/ClickHouse/clickhouse-go/v2"
 	"google.golang.org/grpc/codes"
 )
 
@@ -183,3 +187,41 @@ func PermissionDenied(message string) *Status {
 func PermissionDeniedf(format string, args ...any) *Status {
 	return Newf(codes.PermissionDenied, format, args...)
 }
+
+// ClickHouse error codes for a query that hit a read or time limit.
+const (
+	clickhouseTooManyRows     = 158
+	clickhouseTimeoutExceeded = 159
+	clickhouseTooManyBytes    = 307
+)
+
+var clickhouseCodeRe = regexp.MustCompile(`Code: (\d+)\. DB::Exception`)
+
+// clickhouseCode returns the ClickHouse error code in err, or 0. The HTTP driver returns exceptions as text,
+// the native one as *clickhouse.Exception, so both are read.
+func clickhouseCode(err error) int {
+	var ex *clickhouse.Exception
+	if errors.As(err, &ex) {
+		return int(ex.Code)
+	}
+
+	if m := clickhouseCodeRe.FindStringSubmatch(err.Error()); m != nil {
+		code, _ := strconv.Atoi(m[1])
+
+		return code
+	}
+
+	return 0
+}
+
+// FromQuery maps a failed ClickHouse query to an API status. A query that would read too much or run too
+// long is the caller's to narrow, so it gets a 400 saying how. Anything else is a 500.
+func FromQuery(err error) *Status {
+	switch clickhouseCode(err) {
+	case clickhouseTooManyRows, clickhouseTooManyBytes, clickhouseTimeoutExceeded:
+		return BadRequest("query reads too much data or runs too long: filter the table's primary_key " +
+			"field to a smaller range inside its bounds (see /api/v1/{network}/tables/{table})")
+	default:
+		return Internal(err.Error())
+	}
+}