feat: trace v4 integration

ee/query-service/app/api/api.go

@@ -38,9 +38,10 @@ type APIHandlerOptions struct {
 	Cache                         cache.Cache
 	Gateway                       *httputil.ReverseProxy
 	// Querier Influx Interval
-	FluxInterval     time.Duration
-	UseLogsNewSchema bool
-	UseLicensesV3    bool
+	FluxInterval      time.Duration
+	UseLogsNewSchema  bool
+	UseTraceNewSchema bool
+	UseLicensesV3     bool
 }
 
 type APIHandler struct {
@@ -66,6 +67,7 @@ func NewAPIHandler(opts APIHandlerOptions) (*APIHandler, error) {
 		Cache:                         opts.Cache,
 		FluxInterval:                  opts.FluxInterval,
 		UseLogsNewSchema:              opts.UseLogsNewSchema,
+		UseTraceNewSchema:             opts.UseTraceNewSchema,
 		UseLicensesV3:                 opts.UseLicensesV3,
 	})
 

ee/query-service/app/api/traces.go

@@ -2,32 +2,31 @@ package api
 
 import (
 	"net/http"
-
-	"go.signoz.io/signoz/ee/query-service/app/db"
-	"go.signoz.io/signoz/ee/query-service/model"
-	baseapp "go.signoz.io/signoz/pkg/query-service/app"
-	basemodel "go.signoz.io/signoz/pkg/query-service/model"
-	"go.uber.org/zap"
 )
 
 func (ah *APIHandler) searchTraces(w http.ResponseWriter, r *http.Request) {
 
-	if !ah.CheckFeature(basemodel.SmartTraceDetail) {
-		zap.L().Info("SmartTraceDetail feature is not enabled in this plan")
-		ah.APIHandler.SearchTraces(w, r)
-		return
-	}
-	searchTracesParams, err := baseapp.ParseSearchTracesParams(r)
-	if err != nil {
-		RespondError(w, &model.ApiError{Typ: model.ErrorBadData, Err: err}, "Error reading params")
-		return
-	}
-
-	result, err := ah.opts.DataConnector.SearchTraces(r.Context(), searchTracesParams, db.SmartTraceAlgorithm)
-	if ah.HandleError(w, err, http.StatusBadRequest) {
-		return
-	}
-
-	ah.WriteJSON(w, r, result)
+	ah.APIHandler.SearchTraces(w, r)
+	return
+
+	// This is commented since this will be taken care by new trace API
+
+	// if !ah.CheckFeature(basemodel.SmartTraceDetail) {
+	// 	zap.L().Info("SmartTraceDetail feature is not enabled in this plan")
+	// 	ah.APIHandler.SearchTraces(w, r)
+	// 	return
+	// }
+	// searchTracesParams, err := baseapp.ParseSearchTracesParams(r)
+	// if err != nil {
+	// 	RespondError(w, &model.ApiError{Typ: model.ErrorBadData, Err: err}, "Error reading params")
+	// 	return
+	// }
+
+	// result, err := ah.opts.DataConnector.SearchTraces(r.Context(), searchTracesParams, db.SmartTraceAlgorithm)
+	// if ah.HandleError(w, err, http.StatusBadRequest) {
+	// 	return
+	// }
+
+	// ah.WriteJSON(w, r, result)
 
 }

ee/query-service/app/db/reader.go

@@ -26,8 +26,9 @@ func NewDataConnector(
 	dialTimeout time.Duration,
 	cluster string,
 	useLogsNewSchema bool,
+	useTraceNewSchema bool,
 ) *ClickhouseReader {
-	ch := basechr.NewReader(localDB, promConfigPath, lm, maxIdleConns, maxOpenConns, dialTimeout, cluster, useLogsNewSchema)
+	ch := basechr.NewReader(localDB, promConfigPath, lm, maxIdleConns, maxOpenConns, dialTimeout, cluster, useLogsNewSchema, useTraceNewSchema)
 	return &ClickhouseReader{
 		conn:             ch.GetConn(),
 		appdb:            localDB,

ee/query-service/app/server.go

@@ -77,6 +77,7 @@ type ServerOptions struct {
 	Cluster           string
 	GatewayUrl        string
 	UseLogsNewSchema  bool
+	UseTraceNewSchema bool
 	UseLicensesV3     bool
 }
 
@@ -156,6 +157,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 			serverOptions.DialTimeout,
 			serverOptions.Cluster,
 			serverOptions.UseLogsNewSchema,
+			serverOptions.UseTraceNewSchema,
 		)
 		go qb.Start(readerReady)
 		reader = qb
@@ -189,6 +191,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		serverOptions.DisableRules,
 		lm,
 		serverOptions.UseLogsNewSchema,
+		serverOptions.UseTraceNewSchema,
 	)
 
 	if err != nil {
@@ -270,6 +273,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		FluxInterval:                  fluxInterval,
 		Gateway:                       gatewayProxy,
 		UseLogsNewSchema:              serverOptions.UseLogsNewSchema,
+		UseTraceNewSchema:             serverOptions.UseTraceNewSchema,
 		UseLicensesV3:                 serverOptions.UseLicensesV3,
 	}
 
@@ -737,7 +741,8 @@ func makeRulesManager(
 	cache cache.Cache,
 	disableRules bool,
 	fm baseint.FeatureLookup,
-	useLogsNewSchema bool) (*baserules.Manager, error) {
+	useLogsNewSchema bool,
+	useTraceNewSchema bool) (*baserules.Manager, error) {
 
 	// create engine
 	pqle, err := pqle.FromConfigPath(promConfigPath)
@@ -767,8 +772,9 @@ func makeRulesManager(
 		EvalDelay:    baseconst.GetEvalDelay(),
 
 		PrepareTaskFunc:     rules.PrepareTaskFunc,
-		PrepareTestRuleFunc: rules.TestNotification,
 		UseLogsNewSchema:    useLogsNewSchema,
+		UseTraceNewSchema:   useTraceNewSchema,
+		PrepareTestRuleFunc: rules.TestNotification,
 	}
 
 	// create Manager

ee/query-service/main.go

@@ -94,6 +94,7 @@ func main() {
 	var cluster string
 
 	var useLogsNewSchema bool
+	var useTraceNewSchema bool
 	var useLicensesV3 bool
 	var cacheConfigPath, fluxInterval string
 	var enableQueryServiceLogOTLPExport bool
@@ -105,6 +106,7 @@ func main() {
 	var gatewayUrl string
 
 	flag.BoolVar(&useLogsNewSchema, "use-logs-new-schema", false, "use logs_v2 schema for logs")
+	flag.BoolVar(&useTraceNewSchema, "use-trace-new-schema", false, "use new schema for traces")
 	flag.BoolVar(&useLicensesV3, "use-licenses-v3", false, "use licenses_v3 schema for licenses")
 	flag.StringVar(&promConfigPath, "config", "./config/prometheus.yml", "(prometheus config to read metrics)")
 	flag.StringVar(&skipTopLvlOpsPath, "skip-top-level-ops", "", "(config file to skip top level operations)")
@@ -145,6 +147,7 @@ func main() {
 		Cluster:           cluster,
 		GatewayUrl:        gatewayUrl,
 		UseLogsNewSchema:  useLogsNewSchema,
+		UseTraceNewSchema: useTraceNewSchema,
 		UseLicensesV3:     useLicensesV3,
 	}
 

ee/query-service/rules/manager.go

@@ -26,6 +26,7 @@ func PrepareTaskFunc(opts baserules.PrepareTaskOptions) (baserules.Task, error)
 			opts.FF,
 			opts.Reader,
 			opts.UseLogsNewSchema,
+			opts.UseTraceNewSchema,
 			baserules.WithEvalDelay(opts.ManagerOpts.EvalDelay),
 		)
 
@@ -122,6 +123,7 @@ func TestNotification(opts baserules.PrepareTestRuleOptions) (int, *basemodel.Ap
 			opts.FF,
 			opts.Reader,
 			opts.UseLogsNewSchema,
+			opts.UseTraceNewSchema,
 			baserules.WithSendAlways(),
 			baserules.WithSendUnmatched(),
 		)

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

@@ -167,7 +167,7 @@ func NewReader(
 	dialTimeout time.Duration,
 	cluster string,
 	useLogsNewSchema bool,
-	// useTraceNewSchema bool, // TODO: uncomment this in integration PR
+	useTraceNewSchema bool,
 ) *ClickHouseReader {
 
 	datasource := os.Getenv("ClickHouseUrl")
@@ -178,7 +178,7 @@ func NewReader(
 		zap.L().Fatal("failed to initialize ClickHouse", zap.Error(err))
 	}
 
-	return NewReaderFromClickhouseConnection(db, options, localDB, configFile, featureFlag, cluster, useLogsNewSchema)
+	return NewReaderFromClickhouseConnection(db, options, localDB, configFile, featureFlag, cluster, useLogsNewSchema, useTraceNewSchema)
 }
 
 func NewReaderFromClickhouseConnection(
@@ -189,7 +189,7 @@ func NewReaderFromClickhouseConnection(
 	featureFlag interfaces.FeatureLookup,
 	cluster string,
 	useLogsNewSchema bool,
-	// useTraceNewSchema bool,
+	useTraceNewSchema bool,
 ) *ClickHouseReader {
 	alertManager, err := am.New()
 	if err != nil {
@@ -229,11 +229,10 @@ func NewReaderFromClickhouseConnection(
 
 	traceTableName := options.primary.IndexTable
 	traceLocalTableName := options.primary.LocalIndexTable
-	// TODO: uncomment this in integration PR
-	// if useTraceNewSchema {
-	// 	traceTableName = options.primary.TraceIndexTableV3
-	// 	traceLocalTableName = options.primary.TraceLocalTableNameV3
-	// }
+	if useTraceNewSchema {
+		traceTableName = options.primary.TraceIndexTableV3
+		traceLocalTableName = options.primary.TraceLocalTableNameV3
+	}
 
 	return &ClickHouseReader{
 		db:                      wrap,
@@ -262,7 +261,8 @@ func NewReaderFromClickhouseConnection(
 		cluster:                 cluster,
 		queryProgressTracker:    queryprogress.NewQueryProgressTracker(),
 
-		useLogsNewSchema: useLogsNewSchema,
+		useLogsNewSchema:  useLogsNewSchema,
+		useTraceNewSchema: useTraceNewSchema,
 
 		logsTableV2:              options.primary.LogsTableV2,
 		logsLocalTableV2:         options.primary.LogsLocalTableV2,
@@ -271,7 +271,6 @@ func NewReaderFromClickhouseConnection(
 		logsTableName:            logsTableName,
 		logsLocalTableName:       logsLocalTableName,
 
-		// useTraceNewSchema:    useTraceNewSchema,
 		traceLocalTableName:  traceLocalTableName,
 		traceTableName:       traceTableName,
 		traceResourceTableV3: options.primary.TraceResourceTableV3,

pkg/query-service/app/http_handler.go

@@ -39,6 +39,7 @@ import (
 	querierV2 "go.signoz.io/signoz/pkg/query-service/app/querier/v2"
 	"go.signoz.io/signoz/pkg/query-service/app/queryBuilder"
 	tracesV3 "go.signoz.io/signoz/pkg/query-service/app/traces/v3"
+	tracesV4 "go.signoz.io/signoz/pkg/query-service/app/traces/v4"
 	"go.signoz.io/signoz/pkg/query-service/auth"
 	"go.signoz.io/signoz/pkg/query-service/cache"
 	"go.signoz.io/signoz/pkg/query-service/common"
@@ -110,8 +111,9 @@ type APIHandler struct {
 	// Websocket connection upgrader
 	Upgrader *websocket.Upgrader
 
-	UseLogsNewSchema bool
-	UseLicensesV3    bool
+	UseLogsNewSchema  bool
+	UseTraceNewSchema bool
+	UseLicensesV3     bool
 
 	hostsRepo      *inframetrics.HostsRepo
 	processesRepo  *inframetrics.ProcessesRepo
@@ -163,6 +165,7 @@ type APIHandlerOpts struct {
 	// Use Logs New schema
 	UseLogsNewSchema bool
 
+	UseTraceNewSchema bool
 	// Use Licenses V3 structure
 	UseLicensesV3 bool
 }
@@ -176,21 +179,23 @@ func NewAPIHandler(opts APIHandlerOpts) (*APIHandler, error) {
 	}
 
 	querierOpts := querier.QuerierOptions{
-		Reader:           opts.Reader,
-		Cache:            opts.Cache,
-		KeyGenerator:     queryBuilder.NewKeyGenerator(),
-		FluxInterval:     opts.FluxInterval,
-		FeatureLookup:    opts.FeatureFlags,
-		UseLogsNewSchema: opts.UseLogsNewSchema,
+		Reader:            opts.Reader,
+		Cache:             opts.Cache,
+		KeyGenerator:      queryBuilder.NewKeyGenerator(),
+		FluxInterval:      opts.FluxInterval,
+		FeatureLookup:     opts.FeatureFlags,
+		UseLogsNewSchema:  opts.UseLogsNewSchema,
+		UseTraceNewSchema: opts.UseTraceNewSchema,
 	}
 
 	querierOptsV2 := querierV2.QuerierOptions{
-		Reader:           opts.Reader,
-		Cache:            opts.Cache,
-		KeyGenerator:     queryBuilder.NewKeyGenerator(),
-		FluxInterval:     opts.FluxInterval,
-		FeatureLookup:    opts.FeatureFlags,
-		UseLogsNewSchema: opts.UseLogsNewSchema,
+		Reader:            opts.Reader,
+		Cache:             opts.Cache,
+		KeyGenerator:      queryBuilder.NewKeyGenerator(),
+		FluxInterval:      opts.FluxInterval,
+		FeatureLookup:     opts.FeatureFlags,
+		UseLogsNewSchema:  opts.UseLogsNewSchema,
+		UseTraceNewSchema: opts.UseTraceNewSchema,
 	}
 
 	querier := querier.NewQuerier(querierOpts)
@@ -224,6 +229,7 @@ func NewAPIHandler(opts APIHandlerOpts) (*APIHandler, error) {
 		querier:                       querier,
 		querierV2:                     querierv2,
 		UseLogsNewSchema:              opts.UseLogsNewSchema,
+		UseTraceNewSchema:             opts.UseTraceNewSchema,
 		UseLicensesV3:                 opts.UseLicensesV3,
 		hostsRepo:                     hostsRepo,
 		processesRepo:                 processesRepo,
@@ -242,9 +248,14 @@ func NewAPIHandler(opts APIHandlerOpts) (*APIHandler, error) {
 		logsQueryBuilder = logsv4.PrepareLogsQuery
 	}
 
+	tracesQueryBuilder := tracesV3.PrepareTracesQuery
+	if opts.UseTraceNewSchema {
+		tracesQueryBuilder = tracesV4.PrepareTracesQuery
+	}
+
 	builderOpts := queryBuilder.QueryBuilderOptions{
 		BuildMetricQuery: metricsv3.PrepareMetricQuery,
-		BuildTraceQuery:  tracesV3.PrepareTracesQuery,
+		BuildTraceQuery:  tracesQueryBuilder,
 		BuildLogQuery:    logsQueryBuilder,
 	}
 	aH.queryBuilder = queryBuilder.NewQueryBuilder(builderOpts, aH.featureFlags)
@@ -4330,7 +4341,12 @@ func (aH *APIHandler) queryRangeV3(ctx context.Context, queryRangeParams *v3.Que
 			RespondError(w, apiErrObj, errQuriesByName)
 			return
 		}
-		tracesV3.Enrich(queryRangeParams, spanKeys)
+		if aH.UseTraceNewSchema {
+			tracesV4.Enrich(queryRangeParams, spanKeys)
+		} else {
+			tracesV3.Enrich(queryRangeParams, spanKeys)
+		}
+
 	}
 
 	// WARN: Only works for AND operator in traces query
@@ -4800,7 +4816,11 @@ func (aH *APIHandler) queryRangeV4(ctx context.Context, queryRangeParams *v3.Que
 			RespondError(w, apiErrObj, errQuriesByName)
 			return
 		}
-		tracesV3.Enrich(queryRangeParams, spanKeys)
+		if aH.UseTraceNewSchema {
+			tracesV4.Enrich(queryRangeParams, spanKeys)
+		} else {
+			tracesV3.Enrich(queryRangeParams, spanKeys)
+		}
 	}
 
 	// WARN: Only works for AND operator in traces query