From 9261f848cf30171f1e0f57708f42ea80e1c7bb8f Mon Sep 17 00:00:00 2001
From: Karsten Jeschkies <karsten.jeschkies@grafana.com>
Date: Fri, 10 Nov 2023 20:58:06 +0100
Subject: [PATCH] Serialize unwrap.

---
 pkg/logql/serialize.go      | 62 ++++++++++++++++++++++++++++++++-----
 pkg/logql/serialize_test.go |  9 ++++++
 2 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/pkg/logql/serialize.go b/pkg/logql/serialize.go
index c08db2aa933db..2fa0debc3eede 100644
--- a/pkg/logql/serialize.go
+++ b/pkg/logql/serialize.go
@@ -130,6 +130,8 @@ func (v *JSONSerializer) VisitRangeAggregation(e *syntax.RangeAggregationExpr) {
 	v.WriteMore()
 	v.WriteObjectField("range")
 	v.VisitLogRange(e.Left)
+	v.WriteObjectEnd()
+
 	v.WriteObjectEnd()
 	v.Flush()
 }
@@ -151,8 +153,7 @@ func (v *JSONSerializer) VisitLogRange(e *syntax.LogRange) {
 	if e.Unwrap != nil {
 		v.WriteMore()
 		v.WriteObjectField("unwrap")
-		// TODO:
-		v.WriteObjectEnd()
+		encodeUnwrap(v.Stream, e.Unwrap)
 	}
 
 	v.WriteObjectEnd()
@@ -286,6 +287,47 @@ func decodeGrouping(iter *jsoniter.Iterator) (*syntax.Grouping, error) {
 	return g, nil
 }
 
+func encodeUnwrap(s *jsoniter.Stream, u *syntax.UnwrapExpr) {
+	s.WriteObjectStart()
+	s.WriteObjectField("identifier")
+	s.WriteString(u.Identifier)
+
+	s.WriteMore()
+	s.WriteObjectField("operation")
+	s.WriteString(u.Operation)
+
+	/*
+		s.WriteMore()
+		s.WriteObjectField("post_filterers")
+		s.WriteArrayStart()
+		for i, group := range u.PostFilters{
+			if i > 0 {
+				s.WriteMore()
+			}
+			s.WriteString(group)
+		}
+		s.WriteArrayEnd()
+	*/
+
+	s.WriteObjectEnd()
+}
+
+func decodeUnwrap(iter *jsoniter.Iterator) *syntax.UnwrapExpr {
+	e := &syntax.UnwrapExpr{}
+	for f := iter.ReadObject(); f != ""; f = iter.ReadObject() {
+		switch f {
+		case "identifier":
+			e.Identifier = iter.ReadString()
+		case "operation":
+			e.Operation = iter.ReadString()
+		case "post_filterers":
+			iter.Skip()
+		}
+	}
+
+	return e
+}
+
 func encodeLogSelector(s *jsoniter.Stream, e syntax.LogSelectorExpr) {
 	s.WriteObjectStart()
 	s.WriteObjectField("raw")
@@ -297,6 +339,9 @@ func encodeLogSelector(s *jsoniter.Stream, e syntax.LogSelectorExpr) {
 }
 
 func decodeLogSelector(iter *jsoniter.Iterator) (syntax.LogSelectorExpr, error) {
+	var e syntax.LogSelectorExpr
+	var err error
+
 	for f := iter.ReadObject(); f != ""; f = iter.ReadObject() {
 		switch f {
 		case "raw":
@@ -306,15 +351,16 @@ func decodeLogSelector(iter *jsoniter.Iterator) (syntax.LogSelectorExpr, error)
 				return nil, err
 			}
 
-			if e, ok := expr.(syntax.LogSelectorExpr); ok {
-				return e, nil
-			}
+			var ok bool
+			e, ok = expr.(syntax.LogSelectorExpr)
 
-			return nil, fmt.Errorf("unexpected expression type: want(LogSelectorExpr), got(%T)", expr)
+			if !ok {
+				err = fmt.Errorf("unexpected expression type: want(LogSelectorExpr), got(%T)", expr)
+			}
 		}
 	}
 
-	return nil, fmt.Errorf("missing field: raw")
+	return e, err
 }
 
 func decodeSample(iter *jsoniter.Iterator) (syntax.SampleExpr, error) {
@@ -407,7 +453,7 @@ func decodeLogRange(iter *jsoniter.Iterator) (*syntax.LogRange, error) {
 		case "offset_nanos":
 			expr.Offset = time.Duration(iter.ReadInt64())
 		case "unwrap":
-			iter.Skip() // TODO
+			expr.Unwrap = decodeUnwrap(iter)
 		}
 	}
 
diff --git a/pkg/logql/serialize_test.go b/pkg/logql/serialize_test.go
index 88ba849d5b8b4..498807b8f12a8 100644
--- a/pkg/logql/serialize_test.go
+++ b/pkg/logql/serialize_test.go
@@ -2,6 +2,7 @@ package logql
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -19,6 +20,12 @@ func TestJSONSerializationRoundTrip(t *testing.T) {
 		"simple aggregation": {
 			query: `count_over_time({env="prod", app=~"loki.*"}[5m])`,
 		},
+		"simple aggregation with unwrap": {
+			query: `sum_over_time({env="prod", app=~"loki.*"} | unwrap bytes[5m])`,
+		},
+		"label filterer": {
+			query: `bytes >= 0`,
+		},
 	}
 
 	for name, test := range tests {
@@ -34,6 +41,8 @@ func TestJSONSerializationRoundTrip(t *testing.T) {
 			actual, err := DecodeJSON(buf.String())
 			require.NoError(t, err)
 
+			fmt.Println(buf.String())
+
 			require.Equal(t, test.query, actual.String())
 		})
 	}