fix datasourceRef marshalling when embedded (#981)

Author: ryantxu

experimental/apis/data/v0alpha1/query.go

@@ -16,6 +16,7 @@ import (
 func init() { //nolint:gochecknoinits
 	jsoniter.RegisterTypeEncoder("v0alpha1.DataQuery", &genericQueryCodec{})
 	jsoniter.RegisterTypeDecoder("v0alpha1.DataQuery", &genericQueryCodec{})
+	jsoniter.RegisterTypeDecoder("v0alpha1.DataSourceRef", &datasourceRefCodec{})
 }
 
 type QueryDataRequest struct {
@@ -157,6 +158,7 @@ func (g *DataQuery) GetString(key string) string {
 }
 
 type genericQueryCodec struct{}
+type datasourceRefCodec struct{}
 
 func (codec *genericQueryCodec) IsEmpty(_ unsafe.Pointer) bool {
 	return false
@@ -180,6 +182,33 @@ func (codec *genericQueryCodec) Decode(ptr unsafe.Pointer, iter *j.Iterator) {
 	*((*DataQuery)(ptr)) = q
 }
 
+// Long ago dashboards referenced data sources with only the name
+func (codec *datasourceRefCodec) Decode(ptr unsafe.Pointer, iter *j.Iterator) {
+	q := DataSourceRef{}
+	switch iter.WhatIsNext() {
+	case j.StringValue:
+		q.UID = iter.ReadString()
+	case j.ObjectValue:
+		for field := iter.ReadObject(); field != ""; field = iter.ReadObject() {
+			if iter.Error != nil {
+				return
+			}
+			switch field {
+			case "type":
+				q.Type = iter.ReadString()
+			case "uid":
+				q.UID = iter.ReadString()
+			default:
+				_ = iter.Read() // ignore unused properties
+			}
+		}
+	default:
+		iter.Error = fmt.Errorf("expected string or object")
+		return
+	}
+	*((*DataSourceRef)(ptr)) = q
+}
+
 // MarshalJSON writes JSON including the common and custom values
 func (g DataQuery) MarshalJSON() ([]byte, error) {
 	cfg := j.ConfigCompatibleWithStandardLibrary
@@ -199,6 +228,15 @@ func (g *DataQuery) UnmarshalJSON(b []byte) error {
 	return g.readQuery(iter)
 }
 
+// UnmarshalJSON reads a query from json byte array
+func (g *DataSourceRef) UnmarshalJSON(b []byte) error {
+	iter, err := jsoniter.ParseBytes(jsoniter.ConfigDefault, b)
+	if err != nil {
+		return err
+	}
+	return iter.Unmarshal(b, g)
+}
+
 func (g *DataQuery) DeepCopyInto(out *DataQuery) {
 	*out = *g
 	g.CommonQueryProperties.DeepCopyInto(&out.CommonQueryProperties)
@@ -305,19 +343,7 @@ func (g *CommonQueryProperties) readQuery(iter *jsoniter.Iterator,
 			err = iter.ReadVal(&g.TimeRange)
 		case "datasource":
 			// Old datasource values may just be a string
-			next, err = iter.WhatIsNext()
-			if err != nil {
-				return err
-			}
-			switch next {
-			case j.StringValue:
-				g.Datasource = &DataSourceRef{}
-				g.Datasource.UID, err = iter.ReadString()
-			case j.ObjectValue:
-				err = iter.ReadVal(&g.Datasource)
-			default:
-				return fmt.Errorf("expected string or object")
-			}
+			err = iter.ReadVal(&g.Datasource)
 
 		case "datasourceId":
 			g.DatasourceID, err = iter.ReadInt64()

experimental/apis/data/v0alpha1/query_test.go

@@ -90,6 +90,28 @@ func TestParseQueriesIntoQueryDataRequest(t *testing.T) {
 	})
 }
 
+func TestLegacyDataSourceRef(t *testing.T) {
+	type testWrapper struct {
+		Ref DataSourceRef `json:"ref"`
+	}
+
+	wrap := &testWrapper{}
+	err := json.Unmarshal([]byte(`{ "ref": {"type":"ttt", "uid":"UID"}}`), wrap)
+	require.NoError(t, err)
+	require.Equal(t, "ttt", wrap.Ref.Type)
+	require.Equal(t, "UID", wrap.Ref.UID)
+
+	err = json.Unmarshal([]byte(`{ "ref": "name"}`), wrap)
+	require.NoError(t, err)
+	require.Equal(t, "", wrap.Ref.Type)
+	require.Equal(t, "name", wrap.Ref.UID)
+
+	ref := &DataSourceRef{}
+	err = json.Unmarshal([]byte(`"aaa"`), ref) // string as reference
+	require.NoError(t, err)
+	require.Equal(t, "aaa", ref.UID)
+}
+
 func TestQueryBuilders(t *testing.T) {
 	prop := "testkey"
 	testQ1 := &DataQuery{}