fix(instrumentation): stop recording promql engine spans

Assisted-by: Claude Fable 5

docs/contributing/prometheus.md

@@ -299,8 +299,21 @@ substituted. One subtlety makes it exact: we write stale markers at absent
 grid points. Without them, the engine's lookback would resurrect a point
 from up to `lookback` earlier. The marker encodes "absent here" the way the
 engine itself encodes it. Units evaluate concurrently. Each unit is one
-series lookup plus one grid statement. A step of 0 is an instant query: a
-single evaluation at `end`.
+grid statement: the group-key join resolves the matchers, and the samples
+primary key takes the metric name straight from the selector. Only a
+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`.
+
+Both paths enforce fetch budgets
+(`prometheus::clickhousev2::max_fetched_series` and
+`::max_fetched_samples`; 0 disables). The engine path counts matched series
+and scanned samples. The transpiled path counts buffered grid cells (series
+times grid width) across a plan's units, because transpiled results never
+pass the engine's sample limiter. A refusal is a typed invalid-input error.
+It pierces the engine's `promql.ErrStorage` wrapper
+(`prometheus.TypedStorageError`), so the APIs report a user error, not an
+internal one.
 
 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
@@ -315,8 +328,9 @@ selectors and `last_over_time` transpile at window < step too.
 
 ## Series lookup
 
-Both paths resolve matchers the same way, once per selector
-(`selectSeries`). The series tables hold one row per (fingerprint, bucket)
+The engine path resolves matchers once per selector (`selectSeries`); the
+transpiled path builds the same conditions into its group-key join. Both
+read the same tables. The series tables hold one row per (fingerprint, bucket)
 at 1h/6h/1d/1w granularities. The shared schema package
 (`pkg/telemetryschema/metricstelemetryschema`) picks the table whose bucket
 fits the window. It rounds the window start down to the bucket boundary, so

pkg/instrumentation/promqlspanfilter.go

@@ -0,0 +1,43 @@
+package instrumentation
+
+import (
+	"context"
+	"strings"
+
+	"go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace/embedded"
+	tracenoop "go.opentelemetry.io/otel/trace/noop"
+)
+
+// promqlSpanFilter drops the promql engine's spans. The engine traces every
+// evaluation through the global tracer provider under an unnamed scope: one
+// span per timer plus one per AST node per query ("promqlExec",
+// "promqlInnerEval eval *promql.BinaryExpr", ...), which floods traces
+// without adding value. Filtered spans return a non-recording span that
+// keeps the parent's span context, so descendants (the ClickHouse query
+// spans) still attach to the surrounding span.
+type promqlSpanFilter struct {
+	embedded.TracerProvider
+	delegate trace.TracerProvider
+}
+
+func (p promqlSpanFilter) Tracer(name string, opts ...trace.TracerOption) trace.Tracer {
+	tracer := p.delegate.Tracer(name, opts...)
+	if name != "" {
+		return tracer
+	}
+	return promqlSpanFilterTracer{delegate: tracer, noop: tracenoop.NewTracerProvider().Tracer("")}
+}
+
+type promqlSpanFilterTracer struct {
+	embedded.Tracer
+	delegate trace.Tracer
+	noop     trace.Tracer
+}
+
+func (t promqlSpanFilterTracer) Start(ctx context.Context, spanName string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
+	if strings.HasPrefix(spanName, "promql") {
+		return t.noop.Start(ctx, spanName)
+	}
+	return t.delegate.Start(ctx, spanName, opts...)
+}

pkg/instrumentation/promqlspanfilter_test.go

@@ -0,0 +1,40 @@
+package instrumentation
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+)
+
+func TestPromqlSpanFilter(t *testing.T) {
+	recorder := tracetest.NewSpanRecorder()
+	provider := promqlSpanFilter{delegate: sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))}
+
+	ctx, root := provider.Tracer("http").Start(context.Background(), "GET /api")
+
+	engineCtx, engineSpan := provider.Tracer("").Start(ctx, "promqlInnerEval eval *promql.BinaryExpr")
+	assert.False(t, engineSpan.IsRecording(), "promql engine spans must not record")
+	assert.Equal(t, root.SpanContext().SpanID(), engineSpan.SpanContext().SpanID(), "the filtered span must keep the parent's span context")
+
+	_, child := provider.Tracer("clickhouse").Start(engineCtx, "clickhouse.query")
+	child.End()
+
+	_, other := provider.Tracer("").Start(ctx, "http.request")
+	other.End()
+	root.End()
+
+	var names []string
+	var childParent string
+	for _, span := range recorder.Ended() {
+		names = append(names, span.Name())
+		if span.Name() == "clickhouse.query" {
+			childParent = span.Parent().SpanID().String()
+		}
+	}
+	require.ElementsMatch(t, []string{"clickhouse.query", "http.request", "GET /api"}, names)
+	assert.Equal(t, root.SpanContext().SpanID().String(), childParent, "descendants of a filtered span must attach to the surrounding span")
+}

pkg/instrumentation/sdk.go

@@ -108,7 +108,7 @@ func New(ctx context.Context, cfg Config, build version.Build, serviceName strin
 	}
 
 	// Set the global tracer provider to the sdk tracer provider so that external packages can use this
-	otel.SetTracerProvider(sdk.TracerProvider())
+	otel.SetTracerProvider(promqlSpanFilter{delegate: sdk.TracerProvider()})
 
 	return &SDK{
 		sdk:                       sdk,

pkg/prometheus/clickhouseprometheusv2/capture.go

@@ -73,8 +73,8 @@ func (c *captureQuerier) LabelNames(context.Context, *storage.LabelHints, ...*la
 }
 
 // metricNamesFromMatchers extracts the statically known metric name, if any.
-// The live path derives names from the matched series; the capture path has
-// no execution results, so only a __name__ equality contributes.
+// Only a __name__ equality contributes; a regex selector needs a series
+// lookup to learn the concrete names.
 func metricNamesFromMatchers(matchers []*labels.Matcher) []string {
 	for _, m := range matchers {
 		if m.Name == metricNameLabel && m.Type == labels.MatchEqual && m.Value != "" {

pkg/prometheus/clickhouseprometheusv2/client.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"slices"
 
+	"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 +30,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 +43,7 @@ func newClient(settings factory.ScopedProviderSettings, telemetryStore telemetry
 		settings:       settings,
 		telemetryStore: telemetryStore,
 		lookbackMs:     lookback.Milliseconds(),
+		cfg:            cfg.ClickhouseV2,
 	}
 }
 
@@ -77,6 +80,13 @@ func (c *client) selectSeries(ctx context.Context, query string, args []any) (*s
 		if name := lset.Get(metricNameLabel); name != "" {
 			names[name] = struct{}{}
 		}
+		if c.cfg.MaxFetchedSeries > 0 && len(lookup.fingerprints) > c.cfg.MaxFetchedSeries {
+			return nil, errors.NewInvalidInputf(
+				errors.CodeInvalidInput,
+				"promql selector matched more than %d series; narrow the label matchers or raise prometheus::clickhousev2::max_fetched_series",
+				c.cfg.MaxFetchedSeries,
+			)
+		}
 	}
 	if err := rows.Err(); err != nil {
 		return nil, err
@@ -136,6 +146,8 @@ func (c *client) selectSamples(ctx context.Context, query string, args []any, lo
 		first        = true
 		haveCurrent  bool
 		staleMarker  = math.Float64frombits(promValue.StaleNaN)
+		maxSamples   = c.cfg.MaxFetchedSamples
+		fetched      int64
 		unknownCount int
 	)
 
@@ -144,6 +156,15 @@ func (c *client) selectSamples(ctx context.Context, query string, args []any, lo
 			return nil, err
 		}
 
+		fetched++
+		if maxSamples > 0 && fetched > maxSamples {
+			return nil, errors.NewInvalidInputf(
+				errors.CodeInvalidInput,
+				"promql query would fetch more than %d samples; narrow the selector or time range, or raise prometheus::clickhousev2::max_fetched_samples",
+				maxSamples,
+			)
+		}
+
 		if first || fingerprint != prevFp {
 			first = false
 			prevFp = fingerprint

pkg/prometheus/clickhouseprometheusv2/client_test.go

@@ -0,0 +1,49 @@
+package clickhouseprometheusv2
+
+import (
+	"context"
+	"testing"
+
+	cmock "github.com/SigNoz/clickhouse-go-mock"
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var samplesCols = []cmock.ColumnType{
+	{Name: "fingerprint", Type: "UInt64"},
+	{Name: "unix_milli", Type: "Int64"},
+	{Name: "value", Type: "Float64"},
+	{Name: "flags", Type: "UInt32"},
+}
+
+func TestSelectSeriesBudget(t *testing.T) {
+	c, store := newTestClient(t)
+	c.cfg.MaxFetchedSeries = 1
+
+	store.Mock().ExpectQuery("SELECT fingerprint, any\\(labels\\)").WillReturnRows(cmock.NewRows(seriesCols, [][]any{
+		{uint64(1), `{"__name__":"up","instance":"a"}`},
+		{uint64(2), `{"__name__":"up","instance":"b"}`},
+	}))
+
+	_, err := c.selectSeries(context.Background(), "SELECT fingerprint, any(labels) FROM t", nil)
+	require.Error(t, err)
+	assert.True(t, errors.Ast(err, errors.TypeInvalidInput), "budget refusal must be typed invalid input, got %v", err)
+}
+
+func TestSelectSamplesBudget(t *testing.T) {
+	c, store := newTestClient(t)
+	c.cfg.MaxFetchedSamples = 2
+
+	store.Mock().ExpectQuery("SELECT fingerprint, unix_milli").WillReturnRows(cmock.NewRows(samplesCols, [][]any{
+		{uint64(1), int64(1_700_000_000_000), 1.0, uint32(0)},
+		{uint64(1), int64(1_700_000_060_000), 2.0, uint32(0)},
+		{uint64(1), int64(1_700_000_120_000), 3.0, uint32(0)},
+	}))
+
+	lookup := &seriesLookup{fingerprints: map[uint64]labels.Labels{1: labels.FromStrings("__name__", "up")}}
+	_, err := c.selectSamples(context.Background(), "SELECT fingerprint, unix_milli, value, flags FROM t", nil, lookup)
+	require.Error(t, err)
+	assert.True(t, errors.Ast(err, errors.TypeInvalidInput), "budget refusal must be typed invalid input, got %v", err)
+}

pkg/prometheus/clickhouseprometheusv2/transpiler_exec.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"math"
 	"sort"
+	"sync/atomic"
 	"time"
 
 	"github.com/SigNoz/signoz/pkg/errors"
@@ -88,12 +89,17 @@ func (e *executor) TryExecuteRange(ctx context.Context, qs string, start, end ti
 	}
 
 	// Evaluate every unit concurrently on its own grid (the query grid, or a
-	// subquery grid); each is one series lookup plus one grid query.
+	// subquery grid); each is one grid query (see executeUnit for when a
+	// series lookup precedes it). The units share one grid-cell budget:
+	// transpiled results never pass through the engine's sample limiter, so
+	// without it a large series-count x grid-width query would buffer
+	// unbounded arrays — the OOM this provider exists to prevent.
 	results := make([][]transpiledSeries, len(plan.units))
+	var gridCells atomic.Int64
 	eg, egCtx := errgroup.WithContext(ctx)
 	for i, unit := range plan.units {
 		eg.Go(func() error {
-			res, err := e.executeUnit(egCtx, &unit.core, unit.grid)
+			res, err := e.executeUnit(egCtx, &unit.core, unit.grid, &gridCells)
 			if err != nil {
 				return err
 			}
@@ -133,7 +139,7 @@ type transpiledSeries struct {
 	values []*float64
 }
 
-func (e *executor) executeUnit(ctx context.Context, unit *coreUnit, grid gridContext) ([]transpiledSeries, error) {
+func (e *executor) executeUnit(ctx context.Context, unit *coreUnit, grid gridContext, gridCells *atomic.Int64) ([]transpiledSeries, error) {
 	startMs, endMs, stepMs := grid.startMs, grid.endMs, grid.stepMs
 	windowMs := unit.rangeMs
 	if unit.kind == unitInstant {
@@ -142,19 +148,27 @@ func (e *executor) executeUnit(ctx context.Context, unit *coreUnit, grid gridCon
 	dataStart := startMs - unit.offsetMs - windowMs
 	dataEnd := endMs - unit.offsetMs
 
-	seriesQuery, seriesArgs, err := buildSeriesQuery(dataStart, dataEnd, unit.matchers)
-	if err != nil {
-		return nil, err
-	}
-	lookup, err := e.client.selectSeries(ctx, seriesQuery, seriesArgs)
-	if err != nil {
-		return nil, err
-	}
-	if len(lookup.fingerprints) == 0 {
-		return nil, nil
+	// The group-key join resolves the matchers on its own, so the unit
+	// statement only needs concrete metric names for the samples
+	// primary-key prefix. A selector without a static __name__ learns them
+	// through the series lookup; every other selector skips the roundtrip.
+	metricNames := metricNamesFromMatchers(unit.matchers)
+	if metricNames == nil {
+		seriesQuery, seriesArgs, err := buildSeriesQuery(dataStart, dataEnd, unit.matchers)
+		if err != nil {
+			return nil, err
+		}
+		lookup, err := e.client.selectSeries(ctx, seriesQuery, seriesArgs)
+		if err != nil {
+			return nil, err
+		}
+		if len(lookup.fingerprints) == 0 {
+			return nil, nil
+		}
+		metricNames = lookup.metricNames
 	}
 
-	query, args, err := buildUnitSQL(unit, lookup.metricNames, dataStart, dataEnd, startMs, endMs, stepMs, e.client.lookbackMs)
+	query, args, err := buildUnitSQL(unit, metricNames, dataStart, dataEnd, startMs, endMs, stepMs, e.client.lookbackMs)
 	if err != nil {
 		return nil, err
 	}
@@ -190,6 +204,17 @@ func (e *executor) executeUnit(ctx context.Context, unit *coreUnit, grid gridCon
 		if err := rows.Scan(targets...); err != nil {
 			return nil, err
 		}
+		// One row buffers one grid array; series count times grid width is
+		// the transpiled equivalent of fetched samples. Counted per row as
+		// the arrays accumulate: without a series lookup there is no series
+		// count to charge up front.
+		if maxSamples := e.client.cfg.MaxFetchedSamples; maxSamples > 0 && gridCells.Add(int64(len(gridValues))) > maxSamples {
+			return nil, errors.NewInvalidInputf(
+				errors.CodeInvalidInput,
+				"promql query would buffer more than %d output points; narrow the selector or time range, or raise prometheus::clickhousev2::max_fetched_samples",
+				maxSamples,
+			)
+		}
 		var lset labels.Labels
 		if keyNames != nil {
 			builder := labels.NewScratchBuilder(len(keyNames))

pkg/prometheus/clickhouseprometheusv2/transpiler_test.go

@@ -2,6 +2,7 @@ package clickhouseprometheusv2
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -32,6 +33,17 @@ var seriesCols = []cmock.ColumnType{
 	{Name: "labels", Type: "String"},
 }
 
+var unitCols = []cmock.ColumnType{
+	{Name: "gkey", Type: "String"},
+	{Name: "grid", Type: "Array(Nullable(Float64))"},
+}
+
+// anyArgs matches a bound-argument list by count alone: the mock treats a
+// nil expected argument as a wildcard.
+func anyArgs(n int) []any {
+	return make([]any, n)
+}
+
 func parse(t *testing.T, q string) parser.Expr {
 	t.Helper()
 	expr, err := parser.NewParser(parser.Options{}).ParseExpr(q)
@@ -534,6 +546,30 @@ func TestDisjointWindowLattice(t *testing.T) {
 	}
 }
 
+// Transpiled results never pass the engine's sample limiter, so the grid
+// cells (series x grid width) must be budgeted as the arrays accumulate —
+// otherwise a wide query rebuilds the OOM this provider exists to prevent.
+func TestExecuteUnit_GridCellBudget(t *testing.T) {
+	c, store := newTestClient(t)
+	c.cfg.MaxFetchedSamples = 100
+	e := &executor{client: c, parser: prometheus.NewParser()}
+
+	grid61 := make([]*float64, 61)
+	store.Mock().ExpectQuery("timeSeriesRateToGrid").WithArgs(anyArgs(7)...).WillReturnRows(cmock.NewRows(
+		[]cmock.ColumnType{{Name: "g0", Type: "String"}, {Name: "grid", Type: "Array(Nullable(Float64))"}},
+		[][]any{{"api", grid61}, {"web", grid61}},
+	))
+
+	plan, ok := classify(parse(t, `sum by (job) (rate(up[5m]))`), gridContext{startMs: 1_700_000_000_000, endMs: 1_700_003_600_000, stepMs: 60_000})
+	require.True(t, ok)
+
+	// 2 series x 61 grid points = 122 cells > 100.
+	var cells atomic.Int64
+	_, err := e.executeUnit(context.Background(), &plan.units[0].core, plan.units[0].grid, &cells)
+	require.Error(t, err)
+	assert.True(t, errors.Ast(err, errors.TypeInvalidInput), "budget refusal must be typed invalid input, got %v", err)
+}
+
 func TestTryExecuteRange_WindowedGateFallsBack(t *testing.T) {
 	c, store := newTestClient(t)
 	e := &executor{client: c, parser: prometheus.NewParser()}
@@ -553,7 +589,7 @@ func TestTryExecuteRange_WindowedGateFallsBack(t *testing.T) {
 
 	// 1m range at 5m step: the windows are disjoint slivers — no
 	// divisibility or width requirement, so this transpiles.
-	store.Mock().ExpectQuery("SELECT fingerprint, any\\(labels\\)").WithArgs("up", int64(1_699_999_200_000), int64(1_700_003_600_000)).WillReturnRows(cmock.NewRows(seriesCols, [][]any{}))
+	store.Mock().ExpectQuery("FROM signoz_metrics\\.distributed_samples_v4").WithArgs(anyArgs(9)...).WillReturnRows(cmock.NewRows(unitCols, [][]any{}))
 	_, ok, err = e.TryExecuteRange(context.Background(), `avg_over_time(up[1m])`, start, end, 5*time.Minute)
 	require.NoError(t, err)
 	assert.True(t, ok, "range below step is the disjoint form and must transpile")
@@ -637,12 +673,12 @@ func TestTryExecuteRange_LastStyleWindowBelowStepTranspiles(t *testing.T) {
 	start := time.UnixMilli(1_700_000_000_000)
 	end := time.UnixMilli(1_700_003_600_000)
 
-	store.Mock().ExpectQuery("SELECT fingerprint, any\\(labels\\)").WithArgs("up", int64(1_699_999_200_000), int64(1_700_003_600_000)).WillReturnRows(cmock.NewRows(seriesCols, [][]any{}))
+	store.Mock().ExpectQuery("timeSeriesLastToGrid").WithArgs(anyArgs(10)...).WillReturnRows(cmock.NewRows(unitCols, [][]any{}))
 	_, ok, err := e.TryExecuteRange(context.Background(), `sum by (pod) (up)`, start, end, time.Hour)
 	require.NoError(t, err)
 	assert.True(t, ok, "instant selection at step > lookback must transpile")
 
-	store.Mock().ExpectQuery("SELECT fingerprint, any\\(labels\\)").WithArgs("up", int64(1_699_999_200_000), int64(1_700_003_600_000)).WillReturnRows(cmock.NewRows(seriesCols, [][]any{}))
+	store.Mock().ExpectQuery("timeSeriesLastToGrid").WithArgs(anyArgs(9)...).WillReturnRows(cmock.NewRows(unitCols, [][]any{}))
 	_, ok, err = e.TryExecuteRange(context.Background(), `last_over_time(up[10m])`, start, end, time.Hour)
 	require.NoError(t, err)
 	assert.True(t, ok, "last_over_time at range < step must transpile")

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 (engine path) or buffered grid
+	// cells (transpiled path) one 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,9 @@ type Config struct {
 	// ProviderName selects the storage provider: "clickhouse" (default) or
 	// "clickhousev2".
 	ProviderName string `mapstructure:"provider"`
+
+	// ClickhouseV2 configures the clickhousev2 provider.
+	ClickhouseV2 ClickhouseV2Config `mapstructure:"clickhousev2"`
 }
 
 func NewConfigFactory() factory.ConfigFactory {
@@ -43,6 +56,10 @@ func newConfig() factory.Config {
 		},
 		Timeout:      2 * time.Minute,
 		ProviderName: "clickhouse",
+		ClickhouseV2: ClickhouseV2Config{
+			MaxFetchedSeries:  500_000,
+			MaxFetchedSamples: 50_000_000,
+		},
 	}
 }
 
@@ -53,6 +70,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,22 @@
+package prometheus
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTypedStorageError(t *testing.T) {
+	budget := errors.NewInvalidInputf(errors.CodeInvalidInput, "too many series")
+
+	// The engine wraps a storage error as expanding series: %w inside
+	// promql.ErrStorage, which has no Unwrap method.
+	wrapped := promql.ErrStorage{Err: fmt.Errorf("expanding series: %w", budget)}
+	assert.Equal(t, budget, TypedStorageError(wrapped))
+
+	assert.Nil(t, TypedStorageError(promql.ErrStorage{Err: fmt.Errorf("connection refused")}))
+	assert.Nil(t, TypedStorageError(nil))
+}

pkg/prometheus/promapi/handler.go

@@ -174,6 +174,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 := prometheus.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