From d84fd9493665cb5602b4514f9e746210a32890b8 Mon Sep 17 00:00:00 2001
From: Andre Pereira <adrapereira@gmail.com>
Date: Tue, 6 Aug 2024 18:09:06 +0100
Subject: [PATCH] Tempo: Add migration to enable TraceQL streaming for Tempo
 datasources (#91340)

* Add migration to enable TraceQL streaming for Tempo datasources

* lint

* Always run migration but exit early if feature flag is disabled

* Require feature toggle OR datasource config to enable streaming
---
 .../sqlstore/migrations/migrations.go         |  2 +
 .../migrations/tempo_datasource_mig.go        | 72 +++++++++++++++++++
 .../plugins/datasource/tempo/datasource.ts    |  5 +-
 3 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 pkg/services/sqlstore/migrations/tempo_datasource_mig.go

diff --git a/pkg/services/sqlstore/migrations/migrations.go b/pkg/services/sqlstore/migrations/migrations.go
index 571d47f145a..a3188351898 100644
--- a/pkg/services/sqlstore/migrations/migrations.go
+++ b/pkg/services/sqlstore/migrations/migrations.go
@@ -125,6 +125,8 @@ func (oss *OSSMigrations) AddMigration(mg *Migrator) {
 	ualert.AddRecordingRuleColumns(mg)
 
 	ualert.AddStateResolvedAtColumns(mg)
+
+	enableTraceQLStreaming(mg, oss.features != nil && oss.features.IsEnabledGlobally(featuremgmt.FlagTraceQLStreaming))
 }
 
 func addStarMigrations(mg *Migrator) {
diff --git a/pkg/services/sqlstore/migrations/tempo_datasource_mig.go b/pkg/services/sqlstore/migrations/tempo_datasource_mig.go
new file mode 100644
index 00000000000..df57e7ce0c1
--- /dev/null
+++ b/pkg/services/sqlstore/migrations/tempo_datasource_mig.go
@@ -0,0 +1,72 @@
+package migrations
+
+import (
+	"encoding/json"
+
+	"xorm.io/xorm"
+
+	. "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
+)
+
+func enableTraceQLStreaming(mg *Migrator, enable bool) {
+	mg.AddMigration("Enable traceQL streaming for all Tempo datasources", &AddTraceQLStreamingToJsonData{Enable: enable})
+}
+
+var _ CodeMigration = new(AddTraceQLStreamingToJsonData)
+
+type AddTraceQLStreamingToJsonData struct {
+	MigrationBase
+	Enable bool
+}
+
+func (m *AddTraceQLStreamingToJsonData) SQL(dialect Dialect) string {
+	return "code migration"
+}
+
+type TempoIdJsonDataDTO struct {
+	Id       int64
+	JsonData string
+}
+
+func (m *AddTraceQLStreamingToJsonData) Exec(sess *xorm.Session, mg *Migrator) error {
+	datasources := make([]*TempoIdJsonDataDTO, 0)
+
+	// Skip update if the feature flag is not enabled but mark the migration as completed
+	if !m.Enable {
+		return nil
+	}
+
+	err := sess.SQL("SELECT id, json_data FROM data_source WHERE type = 'tempo'").Find(&datasources)
+
+	if err != nil {
+		return err
+	}
+
+	enabledStreamingMap := map[string]interface{}{
+		"search": true,
+	}
+
+	for _, ds := range datasources {
+		var parsedMap map[string]interface{}
+		if err := json.Unmarshal([]byte(ds.JsonData), &parsedMap); err != nil {
+			continue
+		}
+		// skip datasource if streamingEnabled is already set
+		if parsedMap["streamingEnabled"] != nil {
+			continue
+		}
+		parsedMap["streamingEnabled"] = enabledStreamingMap
+
+		newJsonData, err := json.Marshal(parsedMap)
+		if err != nil {
+			return err
+		}
+
+		_, err = sess.Exec("UPDATE data_source SET json_data = ? WHERE id = ?", newJsonData, ds.Id)
+		if err != nil {
+			return err
+		}
+	}
+
+	return err
+}
diff --git a/public/app/plugins/datasource/tempo/datasource.ts b/public/app/plugins/datasource/tempo/datasource.ts
index 0ccfc56647b..da625c8b8bc 100644
--- a/public/app/plugins/datasource/tempo/datasource.ts
+++ b/public/app/plugins/datasource/tempo/datasource.ts
@@ -277,10 +277,9 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
    */
   isStreamingSearchEnabled() {
     return (
-      config.featureToggles.traceQLStreaming &&
+      (config.featureToggles.traceQLStreaming || this.streamingEnabled?.search) &&
       this.isFeatureAvailable(FeatureName.streaming) &&
-      config.liveEnabled &&
-      this.streamingEnabled?.search
+      config.liveEnabled
     );
   }