Feat: QS: query builder suggestions api v0 (#5634)

* chore: stash initial work with API signature

* chore: put together setup for integration testing filter suggestions

* feat: filter suggestions: suggest attribs using existing autocomplete logic

* chore: filter suggestions test: add expectation for example queries

* feat: filter suggestions: default suggestions when data yet to be received

* feat: finish plumbing basic example queries

* chore: add test for filter suggestions with an existing query

* feat: filter suggestions: don't suggest attribs already included in existing filter

* chore: generate example queries by including existing filter first

* chore: upgrade ClickHouse-go-mock

* chore: some cleanup of reader.GetQBFilterSuggestionsForLogs

* chore: some cleanup of filter suggestion tests

* chore: some cleanup to http handler and request parsing logic for filter suggestions

* chore: remove expectation that attrib suggestions won't contain attribs already used in filter

pkg/query-service/app/clickhouseReader/reader.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"reflect"
 	"regexp"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -4357,6 +4358,128 @@ func (r *ClickHouseReader) GetLogAttributeValues(ctx context.Context, req *v3.Fi
 
 }
 
+func (r *ClickHouseReader) GetQBFilterSuggestionsForLogs(
+	ctx context.Context,
+	req *v3.QBFilterSuggestionsRequest,
+) (*v3.QBFilterSuggestionsResponse, *model.ApiError) {
+	suggestions := v3.QBFilterSuggestionsResponse{
+		AttributeKeys:  []v3.AttributeKey{},
+		ExampleQueries: []v3.FilterSet{},
+	}
+
+	// Use existing autocomplete logic for generating attribute suggestions
+	attribKeysResp, err := r.GetLogAttributeKeys(
+		ctx, &v3.FilterAttributeKeyRequest{
+			SearchText: req.SearchText,
+			DataSource: v3.DataSourceLogs,
+			Limit:      req.Limit,
+		})
+	if err != nil {
+		return nil, model.InternalError(fmt.Errorf("couldn't get attribute keys: %w", err))
+	}
+
+	suggestions.AttributeKeys = attribKeysResp.AttributeKeys
+
+	// Rank suggested attributes
+	slices.SortFunc(suggestions.AttributeKeys, func(a v3.AttributeKey, b v3.AttributeKey) int {
+
+		// Higher score => higher rank
+		attribKeyScore := func(a v3.AttributeKey) int {
+
+			// Scoring criteria is expected to get more sophisticated in follow up changes
+			if a.Type == v3.AttributeKeyTypeResource {
+				return 2
+			}
+
+			if a.Type == v3.AttributeKeyTypeTag {
+				return 1
+			}
+
+			return 0
+		}
+
+		// To sort in descending order of score the return value must be negative when a > b
+		return attribKeyScore(b) - attribKeyScore(a)
+	})
+
+	// Put together suggested example queries.
+
+	newExampleQuery := func() v3.FilterSet {
+		// Include existing filter in example query if specified.
+		if req.ExistingFilter != nil {
+			return *req.ExistingFilter
+		}
+
+		return v3.FilterSet{
+			Operator: "AND",
+			Items:    []v3.FilterItem{},
+		}
+	}
+
+	// Suggest example query for top suggested attribute using existing
+	// autocomplete logic for recommending attrib values
+	//
+	// Example queries for multiple top attributes using a batch version of
+	// GetLogAttributeValues is expected to come in a follow up change
+	if len(suggestions.AttributeKeys) > 0 {
+		topAttrib := suggestions.AttributeKeys[0]
+
+		resp, err := r.GetLogAttributeValues(ctx, &v3.FilterAttributeValueRequest{
+			DataSource:                 v3.DataSourceLogs,
+			FilterAttributeKey:         topAttrib.Key,
+			FilterAttributeKeyDataType: topAttrib.DataType,
+			TagType:                    v3.TagType(topAttrib.Type),
+			Limit:                      1,
+		})
+
+		if err != nil {
+			// Do not fail the entire request if only example query generation fails
+			zap.L().Error("could not find attribute values for creating example query", zap.Error(err))
+
+		} else {
+			addExampleQuerySuggestion := func(value any) {
+				exampleQuery := newExampleQuery()
+
+				exampleQuery.Items = append(exampleQuery.Items, v3.FilterItem{
+					Key:      topAttrib,
+					Operator: "=",
+					Value:    value,
+				})
+
+				suggestions.ExampleQueries = append(
+					suggestions.ExampleQueries, exampleQuery,
+				)
+			}
+
+			if len(resp.StringAttributeValues) > 0 {
+				addExampleQuerySuggestion(resp.StringAttributeValues[0])
+			} else if len(resp.NumberAttributeValues) > 0 {
+				addExampleQuerySuggestion(resp.NumberAttributeValues[0])
+			} else if len(resp.BoolAttributeValues) > 0 {
+				addExampleQuerySuggestion(resp.BoolAttributeValues[0])
+			}
+		}
+	}
+
+	// Suggest static example queries for standard log attributes if needed.
+	if len(suggestions.ExampleQueries) < req.Limit {
+		exampleQuery := newExampleQuery()
+		exampleQuery.Items = append(exampleQuery.Items, v3.FilterItem{
+			Key: v3.AttributeKey{
+				Key:      "body",
+				DataType: v3.AttributeKeyDataTypeString,
+				Type:     v3.AttributeKeyTypeUnspecified,
+				IsColumn: true,
+			},
+			Operator: "contains",
+			Value:    "error",
+		})
+		suggestions.ExampleQueries = append(suggestions.ExampleQueries, exampleQuery)
+	}
+
+	return &suggestions, nil
+}
+
 func readRow(vars []interface{}, columnNames []string, countOfNumberCols int) ([]string, map[string]string, []map[string]string, *v3.Point) {
 	// Each row will have a value and a timestamp, and an optional list of label values
 	// example: {Timestamp: ..., Value: ...}

pkg/query-service/app/http_handler.go

@@ -302,6 +302,8 @@ func (aH *APIHandler) RegisterQueryRangeV3Routes(router *mux.Router, am *AuthMid
 	subRouter.HandleFunc("/query_range", am.ViewAccess(aH.QueryRangeV3)).Methods(http.MethodPost)
 	subRouter.HandleFunc("/query_range/format", am.ViewAccess(aH.QueryRangeV3Format)).Methods(http.MethodPost)
 
+	subRouter.HandleFunc("/filter_suggestions", am.ViewAccess(aH.getQueryBuilderSuggestions)).Methods(http.MethodGet)
+
 	// live logs
 	subRouter.HandleFunc("/logs/livetail", am.ViewAccess(aH.liveTailLogs)).Methods(http.MethodGet)
 }
@@ -3150,6 +3152,30 @@ func (aH *APIHandler) autocompleteAggregateAttributes(w http.ResponseWriter, r *
 	aH.Respond(w, response)
 }
 
+func (aH *APIHandler) getQueryBuilderSuggestions(w http.ResponseWriter, r *http.Request) {
+	req, err := parseQBFilterSuggestionsRequest(r)
+	if err != nil {
+		RespondError(w, err, nil)
+		return
+	}
+
+	if req.DataSource != v3.DataSourceLogs {
+		// Support for traces and metrics might come later
+		RespondError(w, model.BadRequest(
+			fmt.Errorf("suggestions not supported for %s", req.DataSource),
+		), nil)
+		return
+	}
+
+	response, err := aH.reader.GetQBFilterSuggestionsForLogs(r.Context(), req)
+	if err != nil {
+		RespondError(w, err, nil)
+		return
+	}
+
+	aH.Respond(w, response)
+}
+
 func (aH *APIHandler) autoCompleteAttributeKeys(w http.ResponseWriter, r *http.Request) {
 	var response *v3.FilterAttributeKeyResponse
 	req, err := parseFilterAttributeKeyRequest(r)

pkg/query-service/app/parser.go

@@ -2,6 +2,7 @@ package app
 
 import (
 	"bytes"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -837,6 +838,50 @@ func parseAggregateAttributeRequest(r *http.Request) (*v3.AggregateAttributeRequ
 	return &req, nil
 }
 
+func parseQBFilterSuggestionsRequest(r *http.Request) (
+	*v3.QBFilterSuggestionsRequest, *model.ApiError,
+) {
+	dataSource := v3.DataSource(r.URL.Query().Get("dataSource"))
+	if err := dataSource.Validate(); err != nil {
+		return nil, model.BadRequest(err)
+	}
+
+	limit := baseconstants.DefaultFilterSuggestionsLimit
+	limitStr := r.URL.Query().Get("limit")
+	if len(limitStr) > 0 {
+		limit, err := strconv.Atoi(limitStr)
+		if err != nil || limit < 1 {
+			return nil, model.BadRequest(fmt.Errorf(
+				"invalid limit: %s", limitStr,
+			))
+		}
+	}
+
+	var existingFilter *v3.FilterSet
+	existingFilterB64 := r.URL.Query().Get("existingFilter")
+	if len(existingFilterB64) > 0 {
+		decodedFilterJson, err := base64.RawURLEncoding.DecodeString(existingFilterB64)
+		if err != nil {
+			return nil, model.BadRequest(fmt.Errorf("couldn't base64 decode existingFilter: %w", err))
+		}
+
+		existingFilter = &v3.FilterSet{}
+		err = json.Unmarshal(decodedFilterJson, existingFilter)
+		if err != nil {
+			return nil, model.BadRequest(fmt.Errorf("couldn't JSON decode existingFilter: %w", err))
+		}
+	}
+
+	searchText := r.URL.Query().Get("searchText")
+
+	return &v3.QBFilterSuggestionsRequest{
+		DataSource:     dataSource,
+		Limit:          limit,
+		SearchText:     searchText,
+		ExistingFilter: existingFilter,
+	}, nil
+}
+
 func parseFilterAttributeKeyRequest(r *http.Request) (*v3.FilterAttributeKeyRequest, error) {
 	var req v3.FilterAttributeKeyRequest
 

pkg/query-service/constants/constants.go

@@ -407,3 +407,5 @@ var TracesListViewDefaultSelectedColumns = []v3.AttributeKey{
 		IsColumn: true,
 	},
 }
+
+const DefaultFilterSuggestionsLimit = 100

pkg/query-service/interfaces/interface.go

@@ -93,6 +93,10 @@ type Reader interface {
 	GetLogAttributeValues(ctx context.Context, req *v3.FilterAttributeValueRequest) (*v3.FilterAttributeValueResponse, error)
 	GetLogAggregateAttributes(ctx context.Context, req *v3.AggregateAttributeRequest) (*v3.AggregateAttributeResponse, error)
 	GetUsers(ctx context.Context) ([]model.UserPayload, error)
+	GetQBFilterSuggestionsForLogs(
+		ctx context.Context,
+		req *v3.QBFilterSuggestionsRequest,
+	) (*v3.QBFilterSuggestionsResponse, *model.ApiError)
 
 	// Connection needed for rules, not ideal but required
 	GetConn() clickhouse.Conn

pkg/query-service/model/v3/v3.go

@@ -252,6 +252,18 @@ type FilterAttributeKeyRequest struct {
 	Limit              int               `json:"limit"`
 }
 
+type QBFilterSuggestionsRequest struct {
+	DataSource     DataSource `json:"dataSource"`
+	SearchText     string     `json:"searchText"`
+	Limit          int        `json:"limit"`
+	ExistingFilter *FilterSet `json:"existing_filter"`
+}
+
+type QBFilterSuggestionsResponse struct {
+	AttributeKeys  []AttributeKey `json:"attributes"`
+	ExampleQueries []FilterSet    `json:"example_queries"`
+}
+
 type AttributeKeyDataType string
 
 const (

pkg/query-service/tests/integration/filter_suggestions_test.go

@@ -0,0 +1,279 @@
+package tests
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"slices"
+	"strings"
+	"testing"
+
+	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
+	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/pkg/query-service/app"
+	"go.signoz.io/signoz/pkg/query-service/auth"
+	"go.signoz.io/signoz/pkg/query-service/constants"
+	"go.signoz.io/signoz/pkg/query-service/dao"
+	"go.signoz.io/signoz/pkg/query-service/featureManager"
+	"go.signoz.io/signoz/pkg/query-service/model"
+	v3 "go.signoz.io/signoz/pkg/query-service/model/v3"
+	"go.signoz.io/signoz/pkg/query-service/utils"
+)
+
+// If no data has been received yet, filter suggestions should contain
+// standard log fields and static example queries based on them
+func TestDefaultLogsFilterSuggestions(t *testing.T) {
+	require := require.New(t)
+	tb := NewFilterSuggestionsTestBed(t)
+
+	tb.mockAttribKeysQueryResponse([]v3.AttributeKey{})
+	suggestionsQueryParams := map[string]string{}
+	suggestionsResp := tb.GetQBFilterSuggestionsForLogs(suggestionsQueryParams)
+
+	require.Greater(len(suggestionsResp.AttributeKeys), 0)
+	require.True(slices.ContainsFunc(
+		suggestionsResp.AttributeKeys, func(a v3.AttributeKey) bool {
+			return a.Key == "body"
+		},
+	))
+
+	require.Greater(len(suggestionsResp.ExampleQueries), 0)
+	require.False(slices.ContainsFunc(
+		suggestionsResp.AttributeKeys, func(a v3.AttributeKey) bool {
+			return a.Type == v3.AttributeKeyTypeTag || a.Type == v3.AttributeKeyTypeResource
+		},
+	))
+}
+
+func TestLogsFilterSuggestionsWithoutExistingFilter(t *testing.T) {
+	require := require.New(t)
+	tb := NewFilterSuggestionsTestBed(t)
+
+	testAttrib := v3.AttributeKey{
+		Key:      "container_id",
+		Type:     v3.AttributeKeyTypeResource,
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: false,
+	}
+	testAttribValue := "test-container"
+
+	tb.mockAttribKeysQueryResponse([]v3.AttributeKey{testAttrib})
+	tb.mockAttribValuesQueryResponse(testAttrib, []string{testAttribValue})
+	suggestionsQueryParams := map[string]string{}
+	suggestionsResp := tb.GetQBFilterSuggestionsForLogs(suggestionsQueryParams)
+
+	require.Greater(len(suggestionsResp.AttributeKeys), 0)
+	require.True(slices.ContainsFunc(
+		suggestionsResp.AttributeKeys, func(a v3.AttributeKey) bool {
+			return a.Key == testAttrib.Key && a.Type == testAttrib.Type
+		},
+	))
+
+	require.Greater(len(suggestionsResp.ExampleQueries), 0)
+	require.True(slices.ContainsFunc(
+		suggestionsResp.ExampleQueries, func(q v3.FilterSet) bool {
+			return slices.ContainsFunc(q.Items, func(i v3.FilterItem) bool {
+				return i.Key.Key == testAttrib.Key && i.Value == testAttribValue
+			})
+		},
+	))
+}
+
+// If a filter already exists, suggested example queries should
+// contain existing filter
+func TestLogsFilterSuggestionsWithExistingFilter(t *testing.T) {
+	require := require.New(t)
+	tb := NewFilterSuggestionsTestBed(t)
+
+	testAttrib := v3.AttributeKey{
+		Key:      "container_id",
+		Type:     v3.AttributeKeyTypeResource,
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: false,
+	}
+	testAttribValue := "test-container"
+
+	testFilterAttrib := v3.AttributeKey{
+		Key:      "tenant_id",
+		Type:     v3.AttributeKeyTypeTag,
+		DataType: v3.AttributeKeyDataTypeString,
+		IsColumn: false,
+	}
+	testFilterAttribValue := "test-tenant"
+	testFilter := v3.FilterSet{
+		Operator: "AND",
+		Items: []v3.FilterItem{
+			{
+				Key:      testFilterAttrib,
+				Operator: "=",
+				Value:    testFilterAttribValue,
+			},
+		},
+	}
+
+	tb.mockAttribKeysQueryResponse([]v3.AttributeKey{testAttrib, testFilterAttrib})
+	tb.mockAttribValuesQueryResponse(testAttrib, []string{testAttribValue})
+
+	testFilterJson, err := json.Marshal(testFilter)
+	require.Nil(err, "couldn't serialize existing filter to JSON")
+	suggestionsQueryParams := map[string]string{
+		"existingFilter": base64.RawURLEncoding.EncodeToString(testFilterJson),
+	}
+	suggestionsResp := tb.GetQBFilterSuggestionsForLogs(suggestionsQueryParams)
+
+	require.Greater(len(suggestionsResp.AttributeKeys), 0)
+
+	// All example queries should contain the existing filter as a prefix
+	require.Greater(len(suggestionsResp.ExampleQueries), 0)
+	for _, q := range suggestionsResp.ExampleQueries {
+		require.Equal(q.Items[0], testFilter.Items[0])
+	}
+}
+
+// Mocks response for CH queries made by reader.GetLogAttributeKeys
+func (tb *FilterSuggestionsTestBed) mockAttribKeysQueryResponse(
+	attribsToReturn []v3.AttributeKey,
+) {
+	cols := []mockhouse.ColumnType{}
+	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "tagKey"})
+	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "tagType"})
+	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "tagDataType"})
+
+	values := [][]any{}
+	for _, a := range attribsToReturn {
+		rowValues := []any{}
+		rowValues = append(rowValues, a.Key)
+		rowValues = append(rowValues, string(a.Type))
+		rowValues = append(rowValues, string(a.DataType))
+		values = append(values, rowValues)
+	}
+
+	tb.mockClickhouse.ExpectQuery(
+		"select.*from.*signoz_logs.distributed_tag_attributes.*",
+	).WithArgs(
+		constants.DefaultFilterSuggestionsLimit,
+	).WillReturnRows(
+		mockhouse.NewRows(cols, values),
+	)
+
+	// Add expectation for the create table query used to determine
+	// if an attribute is a column
+	cols = []mockhouse.ColumnType{{Type: "String", Name: "statement"}}
+	values = [][]any{{"CREATE TABLE signoz_logs.distributed_logs"}}
+	tb.mockClickhouse.ExpectSelect(
+		"SHOW CREATE TABLE.*",
+	).WillReturnRows(mockhouse.NewRows(cols, values))
+
+}
+
+// Mocks response for CH queries made by reader.GetLogAttributeValues
+func (tb *FilterSuggestionsTestBed) mockAttribValuesQueryResponse(
+	expectedAttrib v3.AttributeKey,
+	stringValuesToReturn []string,
+) {
+	cols := []mockhouse.ColumnType{}
+	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "stringTagValue"})
+
+	values := [][]any{}
+	for _, v := range stringValuesToReturn {
+		rowValues := []any{}
+		rowValues = append(rowValues, v)
+		values = append(values, rowValues)
+	}
+
+	tb.mockClickhouse.ExpectQuery(
+		"select distinct.*stringTagValue.*from.*signoz_logs.distributed_tag_attributes.*",
+	).WithArgs(string(expectedAttrib.Key), v3.TagType(expectedAttrib.Type), 1).WillReturnRows(mockhouse.NewRows(cols, values))
+}
+
+type FilterSuggestionsTestBed struct {
+	t              *testing.T
+	testUser       *model.User
+	qsHttpHandler  http.Handler
+	mockClickhouse mockhouse.ClickConnMockCommon
+}
+
+func (tb *FilterSuggestionsTestBed) GetQBFilterSuggestionsForLogs(
+	queryParams map[string]string,
+) *v3.QBFilterSuggestionsResponse {
+
+	_, dsExistsInQP := queryParams["dataSource"]
+	require.False(tb.t, dsExistsInQP)
+	queryParams["dataSource"] = "logs"
+
+	result := tb.QSGetRequest("/api/v3/filter_suggestions", queryParams)
+
+	dataJson, err := json.Marshal(result.Data)
+	if err != nil {
+		tb.t.Fatalf("could not marshal apiResponse.Data: %v", err)
+	}
+
+	var resp v3.QBFilterSuggestionsResponse
+	err = json.Unmarshal(dataJson, &resp)
+	if err != nil {
+		tb.t.Fatalf("could not unmarshal apiResponse.Data json into PipelinesResponse")
+	}
+
+	return &resp
+}
+
+func NewFilterSuggestionsTestBed(t *testing.T) *FilterSuggestionsTestBed {
+	testDB := utils.NewQueryServiceDBForTests(t)
+
+	fm := featureManager.StartManager()
+	reader, mockClickhouse := NewMockClickhouseReader(t, testDB, fm)
+	mockClickhouse.MatchExpectationsInOrder(false)
+
+	apiHandler, err := app.NewAPIHandler(app.APIHandlerOpts{
+		Reader:       reader,
+		AppDao:       dao.DB(),
+		FeatureFlags: fm,
+	})
+	if err != nil {
+		t.Fatalf("could not create a new ApiHandler: %v", err)
+	}
+
+	router := app.NewRouter()
+	am := app.NewAuthMiddleware(auth.GetUserFromRequest)
+	apiHandler.RegisterRoutes(router, am)
+	apiHandler.RegisterQueryRangeV3Routes(router, am)
+
+	user, apiErr := createTestUser()
+	if apiErr != nil {
+		t.Fatalf("could not create a test user: %v", apiErr)
+	}
+
+	return &FilterSuggestionsTestBed{
+		t:              t,
+		testUser:       user,
+		qsHttpHandler:  router,
+		mockClickhouse: mockClickhouse,
+	}
+}
+
+func (tb *FilterSuggestionsTestBed) QSGetRequest(
+	path string,
+	queryParams map[string]string,
+) *app.ApiResponse {
+	if len(queryParams) > 0 {
+		qps := []string{}
+		for q, v := range queryParams {
+			qps = append(qps, fmt.Sprintf("%s=%s", q, v))
+		}
+		path = fmt.Sprintf("%s?%s", path, strings.Join(qps, "&"))
+	}
+
+	req, err := AuthenticatedRequestForTest(
+		tb.testUser, path, nil,
+	)
+	if err != nil {
+		tb.t.Fatalf("couldn't create authenticated test request: %v", err)
+	}
+
+	result, err := HandleTestRequest(tb.qsHttpHandler, req, 200)
+	if err != nil {
+		tb.t.Fatalf("test request failed: %v", err)
+	}
+	return result
+}

pkg/query-service/tests/integration/logparsingpipeline_test.go

@@ -512,7 +512,7 @@ func (tb *LogPipelinesTestBed) PostPipelinesToQSExpectingStatusCode(
 	postablePipelines logparsingpipeline.PostablePipelines,
 	expectedStatusCode int,
 ) *logparsingpipeline.PipelinesResponse {
-	req, err := NewAuthenticatedTestRequest(
+	req, err := AuthenticatedRequestForTest(
 		tb.testUser, "/api/v1/logs/pipelines", postablePipelines,
 	)
 	if err != nil {
@@ -562,7 +562,7 @@ func (tb *LogPipelinesTestBed) PostPipelinesToQS(
 }
 
 func (tb *LogPipelinesTestBed) GetPipelinesFromQS() *logparsingpipeline.PipelinesResponse {
-	req, err := NewAuthenticatedTestRequest(
+	req, err := AuthenticatedRequestForTest(
 		tb.testUser, "/api/v1/logs/pipelines/latest", nil,
 	)
 	if err != nil {

pkg/query-service/tests/integration/signoz_integrations_test.go

@@ -3,10 +3,7 @@ package tests
 import (
 	"encoding/json"
 	"fmt"
-	"io"
 	"net/http"
-	"net/http/httptest"
-	"runtime/debug"
 	"slices"
 	"testing"
 	"time"
@@ -501,38 +498,18 @@ func (tb *IntegrationsTestBed) RequestQS(
 	path string,
 	postData interface{},
 ) *app.ApiResponse {
-	req, err := NewAuthenticatedTestRequest(
+	req, err := AuthenticatedRequestForTest(
 		tb.testUser, path, postData,
 	)
 	if err != nil {
 		tb.t.Fatalf("couldn't create authenticated test request: %v", err)
 	}
 
-	respWriter := httptest.NewRecorder()
-	tb.qsHttpHandler.ServeHTTP(respWriter, req)
-	response := respWriter.Result()
-	responseBody, err := io.ReadAll(response.Body)
+	result, err := HandleTestRequest(tb.qsHttpHandler, req, 200)
 	if err != nil {
-		tb.t.Fatalf("couldn't read response body received from QS: %v", err)
+		tb.t.Fatalf("test request failed: %v", err)
 	}
-
-	if response.StatusCode != 200 {
-		tb.t.Fatalf(
-			"unexpected response status from query service for path %s. status: %d, body: %v\n%v",
-			path, response.StatusCode, string(responseBody), string(debug.Stack()),
-		)
-	}
-
-	var result app.ApiResponse
-	err = json.Unmarshal(responseBody, &result)
-	if err != nil {
-		tb.t.Fatalf(
-			"Could not unmarshal QS response into an ApiResponse.\nResponse body: %s",
-			string(responseBody),
-		)
-	}
-
-	return &result
+	return result
 }
 
 func (tb *IntegrationsTestBed) mockLogQueryResponse(logsInResponse []model.SignozLog) {

pkg/query-service/tests/integration/test_utils.go

@@ -5,8 +5,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"runtime/debug"
 	"testing"
 	"time"
 
@@ -16,6 +18,7 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/entry"
 	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
 	"github.com/stretchr/testify/require"
+	"go.signoz.io/signoz/pkg/query-service/app"
 	"go.signoz.io/signoz/pkg/query-service/app/clickhouseReader"
 	"go.signoz.io/signoz/pkg/query-service/auth"
 	"go.signoz.io/signoz/pkg/query-service/constants"
@@ -172,7 +175,7 @@ func createTestUser() (*model.User, *model.ApiError) {
 	)
 }
 
-func NewAuthenticatedTestRequest(
+func AuthenticatedRequestForTest(
 	user *model.User,
 	path string,
 	postData interface{},
@@ -198,3 +201,31 @@ func NewAuthenticatedTestRequest(
 	req.Header.Add("Authorization", "Bearer "+userJwt.AccessJwt)
 	return req, nil
 }
+
+func HandleTestRequest(handler http.Handler, req *http.Request, expectedStatus int) (*app.ApiResponse, error) {
+	respWriter := httptest.NewRecorder()
+	handler.ServeHTTP(respWriter, req)
+	response := respWriter.Result()
+	responseBody, err := io.ReadAll(response.Body)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't read response body received from QS: %w", err)
+	}
+
+	if response.StatusCode != expectedStatus {
+		return nil, fmt.Errorf(
+			"unexpected response status from query service for path %s. status: %d, body: %v\n%v",
+			req.URL.Path, response.StatusCode, string(responseBody), string(debug.Stack()),
+		)
+	}
+
+	var result app.ApiResponse
+	err = json.Unmarshal(responseBody, &result)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"Could not unmarshal QS response into an ApiResponse.\nResponse body: %s",
+			string(responseBody),
+		)
+	}
+
+	return &result, nil
+}