From 834d70f52b6e1046b47933e7247b2e63c3d71263 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=A9vin=20Gomez?= <kevin.gomez@grafana.com>
Date: Fri, 27 Sep 2024 22:44:27 +0200
Subject: [PATCH] Use the datasource type within queries to unmarshal them

---
 .../golang/templates/runtime/runtime.tmpl     | 21 +++++++++++++++++++
 testdata/generated/cog/runtime.go             | 21 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/internal/jennies/golang/templates/runtime/runtime.tmpl b/internal/jennies/golang/templates/runtime/runtime.tmpl
index b3520cc2..1bfa3f92 100644
--- a/internal/jennies/golang/templates/runtime/runtime.tmpl
+++ b/internal/jennies/golang/templates/runtime/runtime.tmpl
@@ -69,6 +69,27 @@ func (runtime *Runtime) UnmarshalDataquery(raw []byte, dataqueryTypeHint string)
 		}
 	}
 
+	// Dataqueries might reference the datasource to use, and its type. Let's use that.
+	partialDataquery := struct {
+		Datasource struct {
+			Type string `json:"type"`
+		} `json:"datasource"`
+	}{}
+	if err := json.Unmarshal(raw, &partialDataquery); err != nil {
+		return nil, err
+	}
+	if partialDataquery.Datasource.Type != "" {
+		config, found := runtime.dataqueryVariants[partialDataquery.Datasource.Type]
+		if found {
+			dataquery, err := config.DataqueryUnmarshaler(raw)
+			if err != nil {
+				return nil, err
+			}
+
+			return dataquery.(variants.Dataquery), nil
+		}
+	}
+
 	// We have no idea what type the dataquery is: use our `UnknownDataquery` bag to not lose data.
 	dataquery := variants.UnknownDataquery{}
 	if err := json.Unmarshal(raw, &dataquery); err != nil {
diff --git a/testdata/generated/cog/runtime.go b/testdata/generated/cog/runtime.go
index 0777ccf2..21389160 100644
--- a/testdata/generated/cog/runtime.go
+++ b/testdata/generated/cog/runtime.go
@@ -78,6 +78,27 @@ func (runtime *Runtime) UnmarshalDataquery(raw []byte, dataqueryTypeHint string)
 		}
 	}
 
+	// Dataqueries might reference the datasource to use, and its type. Let's use that.
+	partialDataquery := struct {
+		Datasource struct {
+			Type string `json:"type"`
+		} `json:"datasource"`
+	}{}
+	if err := json.Unmarshal(raw, &partialDataquery); err != nil {
+		return nil, err
+	}
+	if partialDataquery.Datasource.Type != "" {
+		config, found := runtime.dataqueryVariants[partialDataquery.Datasource.Type]
+		if found {
+			dataquery, err := config.DataqueryUnmarshaler(raw)
+			if err != nil {
+				return nil, err
+			}
+
+			return dataquery.(variants.Dataquery), nil
+		}
+	}
+
 	// We have no idea what type the dataquery is: use our `UnknownDataquery` bag to not lose data.
 	dataquery := variants.UnknownDataquery{}
 	if err := json.Unmarshal(raw, &dataquery); err != nil {