feat(clickhouseprometheusv2): enforce fetch budgets in clickhouse on the engine read path

Assisted-by: Claude Fable 5

.github/workflows/integrationci.yaml

@@ -59,6 +59,7 @@ jobs:
           - rawexportdata
           - promqlconformance
           - promapiconformance
+          - promqlbudget
           - querierauthz
           - role
           - rootuser

docs/contributing/prometheus.md

@@ -305,6 +305,18 @@ selector without a static `__name__` runs the series lookup first, to learn
 the concrete metric names. A step of 0 is an instant query: a single
 evaluation at `end`.
 
+The engine path enforces fetch budgets in ClickHouse
+(`prometheus::clickhousev2::max_fetched_series` and
+`::max_fetched_samples`; 0 disables). The series lookup and the samples
+query carry `max_result_rows` with `result_overflow_mode = 'throw'`, so an
+over-budget query stops in the database instead of streaming into the
+service. The client maps the refusal (`TOO_MANY_ROWS_OR_BYTES`) to a typed
+invalid-input error. The error pierces the
+engine's `promql.ErrStorage` wrapper (`prometheus.TypedStorageError`), so
+the APIs report a user error, not an internal one. Transpiled statements
+carry no result budget: their result rows are output series, which the
+querier fleet caps by other means.
+
 A note on the window sliver: when the window is narrower than the step, the
 grid windows cover only `window/step` of the timeline. A sample in a gap
 belongs to no window. It cannot move any grid point, but the grid aggregate

pkg/prometheus/clickhouseprometheusv2/client.go

@@ -3,10 +3,15 @@ package clickhouseprometheusv2
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"log/slog"
 	"math"
 	"slices"
 
+	chproto "github.com/ClickHouse/ch-go/proto"
+	"github.com/ClickHouse/clickhouse-go/v2"
+
+	"github.com/SigNoz/signoz/pkg/errors"
 	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/prometheus"
 	"github.com/SigNoz/signoz/pkg/telemetrystore"
@@ -29,6 +34,7 @@ type client struct {
 	settings       factory.ScopedProviderSettings
 	telemetryStore telemetrystore.TelemetryStore
 	lookbackMs     int64
+	cfg            prometheus.ClickhouseV2Config
 }
 
 func newClient(settings factory.ScopedProviderSettings, telemetryStore telemetrystore.TelemetryStore, cfg prometheus.Config) *client {
@@ -41,6 +47,7 @@ func newClient(settings factory.ScopedProviderSettings, telemetryStore telemetry
 		settings:       settings,
 		telemetryStore: telemetryStore,
 		lookbackMs:     lookback.Milliseconds(),
+		cfg:            cfg.ClickhouseV2,
 	}
 }
 
@@ -52,11 +59,49 @@ func (c *client) withContext(ctx context.Context, functionName string) context.C
 	})
 }
 
+// budgetSettings caps the result server-side; 'throw' refuses the query
+// instead of silently truncating the result.
+func budgetSettings(maxRows int64) string {
+	return fmt.Sprintf(" SETTINGS max_result_rows = %d, result_overflow_mode = 'throw'", maxRows)
+}
+
+// budgetExceeded reports whether err is ClickHouse refusing a query over its
+// result limit (the budgets above).
+func budgetExceeded(err error) bool {
+	var chErr *clickhouse.Exception
+	return errors.As(err, &chErr) && chErr.Code == int32(chproto.ErrTooManyRowsOrBytes)
+}
+
+func (c *client) seriesError(err error) error {
+	if budgetExceeded(err) {
+		return errors.NewInvalidInputf(
+			errors.CodeInvalidInput,
+			"promql selector matched more than %d series; narrow the label matchers",
+			c.cfg.MaxFetchedSeries,
+		)
+	}
+	return err
+}
+
+func (c *client) samplesError(err error) error {
+	if budgetExceeded(err) {
+		return errors.NewInvalidInputf(
+			errors.CodeInvalidInput,
+			"promql query would fetch more than %d samples; narrow the selector or time range",
+			c.cfg.MaxFetchedSamples,
+		)
+	}
+	return err
+}
+
 func (c *client) selectSeries(ctx context.Context, query string, args []any) (*seriesLookup, error) {
 	ctx = c.withContext(ctx, "selectSeries")
+	if c.cfg.MaxFetchedSeries > 0 {
+		query += budgetSettings(int64(c.cfg.MaxFetchedSeries))
+	}
 	rows, err := c.telemetryStore.ClickhouseDB().Query(ctx, query, args...)
 	if err != nil {
-		return nil, err
+		return nil, c.seriesError(err)
 	}
 	defer rows.Close()
 
@@ -79,7 +124,7 @@ func (c *client) selectSeries(ctx context.Context, query string, args []any) (*s
 		}
 	}
 	if err := rows.Err(); err != nil {
-		return nil, err
+		return nil, c.seriesError(err)
 	}
 
 	for name := range names {
@@ -119,9 +164,12 @@ func unmarshalLabels(s string) (labels.Labels, error) {
 // the engine as-is over the same dirty data.
 func (c *client) selectSamples(ctx context.Context, query string, args []any, lookup *seriesLookup) ([]*series, error) {
 	ctx = c.withContext(ctx, "selectSamples")
+	if c.cfg.MaxFetchedSamples > 0 {
+		query += budgetSettings(c.cfg.MaxFetchedSamples)
+	}
 	rows, err := c.telemetryStore.ClickhouseDB().Query(ctx, query, args...)
 	if err != nil {
-		return nil, err
+		return nil, c.samplesError(err)
 	}
 	defer rows.Close()
 
@@ -168,7 +216,7 @@ func (c *client) selectSamples(ctx context.Context, query string, args []any, lo
 		current.vs = append(current.vs, val)
 	}
 	if err := rows.Err(); err != nil {
-		return nil, err
+		return nil, c.samplesError(err)
 	}
 
 	if unknownCount > 0 {

pkg/prometheus/config.go

@@ -13,6 +13,16 @@ type ActiveQueryTrackerConfig struct {
 	MaxConcurrent int    `mapstructure:"max_concurrent"`
 }
 
+type ClickhouseV2Config struct {
+	// MaxFetchedSeries caps the series one selector may match; 0 disables
+	// the cap.
+	MaxFetchedSeries int `mapstructure:"max_fetched_series"`
+
+	// MaxFetchedSamples caps the samples one engine-path query may fetch;
+	// 0 disables the cap.
+	MaxFetchedSamples int64 `mapstructure:"max_fetched_samples"`
+}
+
 type Config struct {
 	ActiveQueryTrackerConfig ActiveQueryTrackerConfig `mapstructure:"active_query_tracker"`
 
@@ -28,6 +38,8 @@ type Config struct {
 	// ProviderName selects the storage provider: "clickhouse" (default) or
 	// "clickhousev2".
 	ProviderName string `mapstructure:"provider"`
+
+	ClickhouseV2 ClickhouseV2Config `mapstructure:"clickhousev2"`
 }
 
 func NewConfigFactory() factory.ConfigFactory {
@@ -43,6 +55,10 @@ func newConfig() factory.Config {
 		},
 		Timeout:      2 * time.Minute,
 		ProviderName: "clickhouse",
+		ClickhouseV2: ClickhouseV2Config{
+			MaxFetchedSeries:  500_000,
+			MaxFetchedSamples: 50_000_000,
+		},
 	}
 }
 
@@ -53,6 +69,9 @@ func (c Config) Validate() error {
 	if c.ProviderName != "" && c.ProviderName != "clickhouse" && c.ProviderName != "clickhousev2" {
 		return errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput, "prometheus::provider must be one of [clickhouse, clickhousev2], got %q", c.ProviderName)
 	}
+	if c.ClickhouseV2.MaxFetchedSeries < 0 || c.ClickhouseV2.MaxFetchedSamples < 0 {
+		return errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput, "prometheus::clickhousev2 limits must not be negative")
+	}
 	return nil
 }
 

pkg/prometheus/errors.go

@@ -0,0 +1,30 @@
+package prometheus
+
+import (
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/prometheus/prometheus/promql"
+)
+
+// TypedStorageError walks an engine execution error chain looking for a
+// SigNoz-typed invalid-input error raised by the storage layer (the fetch
+// budget refusals). Every wrapper level is stepped through by hand: Ast is a
+// bare type cast, not an unwrap — it misses a typed error behind the
+// engine's "expanding series: %w" — and promql.ErrStorage has no Unwrap
+// method at all, so a plain unwrap loop would stop at it.
+func TypedStorageError(execErr error) error {
+	for e := execErr; e != nil; {
+		if errors.Ast(e, errors.TypeInvalidInput) {
+			return e
+		}
+		if es, ok := e.(promql.ErrStorage); ok {
+			e = es.Err
+			continue
+		}
+		u, ok := e.(interface{ Unwrap() error })
+		if !ok {
+			return nil
+		}
+		e = u.Unwrap()
+	}
+	return nil
+}

pkg/prometheus/errors_test.go

@@ -0,0 +1,28 @@
+package prometheus
+
+import (
+	"testing"
+
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/stretchr/testify/assert"
+)
+
+// expandWrap mirrors the engine's "expanding series: %w" wrapper: a plain
+// message-carrying layer with a single Unwrap.
+type expandWrap struct{ inner error }
+
+func (w expandWrap) Error() string { return "expanding series: " + w.inner.Error() }
+func (w expandWrap) Unwrap() error { return w.inner }
+
+func TestTypedStorageError(t *testing.T) {
+	budget := errors.NewInvalidInputf(errors.CodeInvalidInput, "too many series")
+
+	// promql.ErrStorage has no Unwrap method; the walk must pierce it by
+	// type and then step through the plain wrapper.
+	wrapped := promql.ErrStorage{Err: expandWrap{inner: budget}}
+	assert.Equal(t, budget, TypedStorageError(wrapped))
+
+	assert.Nil(t, TypedStorageError(promql.ErrStorage{Err: errors.Newf(errors.TypeInternal, errors.CodeInternal, "connection refused")}))
+	assert.Nil(t, TypedStorageError(nil))
+}

pkg/prometheus/handler.go

@@ -153,6 +153,13 @@ func (h *handler) exec(ctx context.Context, w http.ResponseWriter, r *http.Reque
 		case promql.ErrQueryTimeout:
 			h.respondError(ctx, w, errTimeout, res.Err)
 		case promql.ErrStorage:
+			// A fetch-budget refusal is the storage-level twin of the
+			// engine's own too-many-samples error, which upstream maps to
+			// "execution", not "internal".
+			if typed := TypedStorageError(res.Err); typed != nil {
+				h.respondError(ctx, w, errExec, typed)
+				return
+			}
 			h.respondError(ctx, w, errInternal, res.Err)
 		default:
 			h.respondError(ctx, w, errExec, res.Err)

pkg/querier/promql_query.go

@@ -43,6 +43,10 @@ var quotedMetricOutsideBracesPattern = regexp.MustCompile(`"([^"]+)"\s*\{`)
 // tryEnhancePromQLExecError attempts to convert a PromQL execution error into
 // a properly typed error. Returns nil if the error is not a recognized execution error.
 func tryEnhancePromQLExecError(execErr error) error {
+	if typed := prometheus.TypedStorageError(execErr); typed != nil {
+		return typed
+	}
+
 	var eqc promql.ErrQueryCanceled
 	var eqt promql.ErrQueryTimeout
 	var es promql.ErrStorage

tests/integration/tests/promqlbudget/01_budgets.py

@@ -0,0 +1,93 @@
+from collections.abc import Callable
+from datetime import UTC, datetime, timedelta
+from http import HTTPStatus
+
+import requests
+
+from fixtures import types
+from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD
+from fixtures.metrics import Metrics
+
+QUERY_TIMEOUT = 30
+
+
+def test_fetch_budgets(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_metrics: Callable[[list[Metrics]], None],
+) -> None:
+    # The container runs with max_fetched_series=10 and
+    # max_fetched_samples=100 (see conftest.py). Both metrics take the
+    # engine path: instant queries always do, and changes() is not
+    # transpilable, so the range query falls back too.
+    end = datetime.now(tz=UTC).replace(second=0, microsecond=0) - timedelta(minutes=5)
+    metrics = [
+        Metrics(
+            metric_name="budget_series_metric",
+            labels={"instance": str(i)},
+            timestamp=end,
+            value=1.0,
+            flags=0,
+        )
+        for i in range(25)
+    ]
+    start = end - timedelta(minutes=50)
+    metrics.extend(
+        Metrics(
+            metric_name="budget_samples_metric",
+            labels={"instance": "0"},
+            timestamp=start + timedelta(seconds=15 * i),
+            value=float(i),
+            flags=0,
+        )
+        for i in range(200)
+    )
+    insert_metrics(metrics)
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    headers = {"authorization": f"Bearer {token}"}
+
+    # 25 matched series exceed the series budget of 10: the series lookup is
+    # refused by ClickHouse and surfaces as a Prometheus execution error.
+    response = requests.get(
+        signoz.self.host_configs["8080"].get("/prometheus/api/v1/query"),
+        params={"query": "budget_series_metric", "time": end.timestamp()},
+        timeout=QUERY_TIMEOUT,
+        headers=headers,
+    )
+    assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
+    body = response.json()
+    assert body["status"] == "error"
+    assert body["errorType"] == "execution"
+    assert "matched more than 10 series" in body["error"]
+
+    # The range selector fetches all 200 raw samples, exceeding the samples
+    # budget of 100.
+    response = requests.get(
+        signoz.self.host_configs["8080"].get("/prometheus/api/v1/query_range"),
+        params={
+            "query": "changes(budget_samples_metric[30m])",
+            "start": start.timestamp(),
+            "end": end.timestamp(),
+            "step": 60,
+        },
+        timeout=QUERY_TIMEOUT,
+        headers=headers,
+    )
+    assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
+    body = response.json()
+    assert body["status"] == "error"
+    assert body["errorType"] == "execution"
+    assert "more than 100 samples" in body["error"]
+
+    # One series and one last-sample fetch sit inside both budgets.
+    response = requests.get(
+        signoz.self.host_configs["8080"].get("/prometheus/api/v1/query"),
+        params={"query": "budget_samples_metric", "time": end.timestamp()},
+        timeout=QUERY_TIMEOUT,
+        headers=headers,
+    )
+    assert response.status_code == HTTPStatus.OK
+    body = response.json()
+    assert body["status"] == "success"
+    assert len(body["data"]["result"]) == 1

tests/integration/tests/promqlbudget/conftest.py

@@ -0,0 +1,37 @@
+import pytest
+from testcontainers.core.container import Network
+
+from fixtures import types
+from fixtures.signoz import create_signoz
+
+
+@pytest.fixture(name="signoz", scope="package")
+def signoz_promql_budget(
+    network: Network,
+    migrator: types.Operation,  # pylint: disable=unused-argument
+    zeus: types.TestContainerDocker,
+    gateway: types.TestContainerDocker,
+    sqlstore: types.TestContainerSQL,
+    clickhouse: types.TestContainerClickhouse,
+    request: pytest.FixtureRequest,
+    pytestconfig: pytest.Config,
+) -> types.SigNoz:
+    """
+    SigNoz with clickhousev2 serving and tiny fetch budgets, so engine-path
+    queries trip the ClickHouse-enforced result limits with small fixtures.
+    """
+    return create_signoz(
+        network=network,
+        zeus=zeus,
+        gateway=gateway,
+        sqlstore=sqlstore,
+        clickhouse=clickhouse,
+        request=request,
+        pytestconfig=pytestconfig,
+        cache_key="signoz-promql-budget",
+        env_overrides={
+            "SIGNOZ_PROMETHEUS_PROVIDER": "clickhousev2",
+            "SIGNOZ_PROMETHEUS_CLICKHOUSEV2_MAX__FETCHED__SERIES": "10",
+            "SIGNOZ_PROMETHEUS_CLICKHOUSEV2_MAX__FETCHED__SAMPLES": "100",
+        },
+    )