From 24a5d8dde69b664230e2a74ed00447876f9f956c Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Wed, 5 Jun 2024 16:14:23 +0300
Subject: [PATCH] feat(chsql): add `ORDER BY` expression

---
 internal/chstorage/chsql/_golden/Test2.sql |  2 +-
 internal/chstorage/chsql/_golden/Test3.sql |  1 +
 internal/chstorage/chsql/_golden/Test4.sql |  1 +
 internal/chstorage/chsql/chsql.go          | 41 ++++++++++++++++++++++
 internal/chstorage/chsql/select.go         | 33 +++++++++++++++++
 internal/chstorage/chsql/select_test.go    | 28 +++++++++++++++
 6 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 internal/chstorage/chsql/_golden/Test3.sql
 create mode 100644 internal/chstorage/chsql/_golden/Test4.sql

diff --git a/internal/chstorage/chsql/_golden/Test2.sql b/internal/chstorage/chsql/_golden/Test2.sql
index 3f36af74..a86886cf 100644
--- a/internal/chstorage/chsql/_golden/Test2.sql
+++ b/internal/chstorage/chsql/_golden/Test2.sql
@@ -1 +1 @@
-SELECT (span_id) FROM (SELECT (span_id),(timestamp) FROM spans WHERE (true) AND (duration > 3.14) AND (duration < 3.14))
\ No newline at end of file
+SELECT (span_id) FROM (SELECT (span_id),(timestamp) FROM spans WHERE (true) AND (duration > 3.14) AND (duration < 3.14)) LIMIT 1
\ No newline at end of file
diff --git a/internal/chstorage/chsql/_golden/Test3.sql b/internal/chstorage/chsql/_golden/Test3.sql
new file mode 100644
index 00000000..2d2a811c
--- /dev/null
+++ b/internal/chstorage/chsql/_golden/Test3.sql
@@ -0,0 +1 @@
+SELECT (timestamp) FROM spans ORDER BY timestamp DESC
\ No newline at end of file
diff --git a/internal/chstorage/chsql/_golden/Test4.sql b/internal/chstorage/chsql/_golden/Test4.sql
new file mode 100644
index 00000000..2abe069d
--- /dev/null
+++ b/internal/chstorage/chsql/_golden/Test4.sql
@@ -0,0 +1 @@
+SELECT (timestamp) FROM spans ORDER BY timestamp ASC,duration DESC
\ No newline at end of file
diff --git a/internal/chstorage/chsql/chsql.go b/internal/chstorage/chsql/chsql.go
index 7baec4b5..b81aa505 100644
--- a/internal/chstorage/chsql/chsql.go
+++ b/internal/chstorage/chsql/chsql.go
@@ -2,11 +2,32 @@
 package chsql
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/go-faster/errors"
 )
 
+// Order defines sorting order.
+type Order uint8
+
+// String implements [fmt.Stringer].
+func (o Order) String() string {
+	switch o {
+	case Desc:
+		return "DESC"
+	case Asc:
+		return "ASC"
+	default:
+		return fmt.Sprintf("unknown order %d", o)
+	}
+}
+
+const (
+	Desc Order = iota + 1
+	Asc
+)
+
 // Printer prints SQL query.
 type Printer struct {
 	sb        strings.Builder
@@ -164,6 +185,26 @@ func (p *Printer) Where() {
 	p.Ident("WHERE")
 }
 
+// Order writes `ORDER` ident.
+func (p *Printer) Order() {
+	p.Ident("ORDER")
+}
+
+// By writes `BY` ident.
+func (p *Printer) By() {
+	p.Ident("BY")
+}
+
+// Asc writes `ASC` ident.
+func (p *Printer) Asc() {
+	p.Ident("ASC")
+}
+
+// Desc writes `DESC` ident.
+func (p *Printer) Desc() {
+	p.Ident("DESC")
+}
+
 // Limit writes `LIMIT` ident.
 func (p *Printer) Limit() {
 	p.Ident("LIMIT")
diff --git a/internal/chstorage/chsql/select.go b/internal/chstorage/chsql/select.go
index 1af740c9..88ec6246 100644
--- a/internal/chstorage/chsql/select.go
+++ b/internal/chstorage/chsql/select.go
@@ -18,10 +18,16 @@ type SelectQuery struct {
 
 	// where is a set of expression joined by AND
 	where []expr
+	order []orderExpr
 
 	limit int
 }
 
+type orderExpr struct {
+	expr  expr
+	order Order
+}
+
 // Select creates a new [SelectQuery].
 func Select(table string, columns ...ResultColumn) *SelectQuery {
 	return &SelectQuery{
@@ -50,6 +56,12 @@ func (q *SelectQuery) Where(filters ...expr) *SelectQuery {
 	return q
 }
 
+// Order adds order to query.
+func (q *SelectQuery) Order(e expr, order Order) *SelectQuery {
+	q.order = append(q.order, orderExpr{expr: e, order: order})
+	return q
+}
+
 // Limit sets query limit.
 //
 // If n is equal to or less than zero, limit is ignored.
@@ -109,6 +121,27 @@ func (q *SelectQuery) WriteSQL(p *Printer) error {
 			p.CloseParen()
 		}
 	}
+	if len(q.order) > 0 {
+		p.Order()
+		p.By()
+
+		for i, e := range q.order {
+			if i != 0 {
+				p.Comma()
+			}
+			if err := p.WriteExpr(e.expr); err != nil {
+				return errors.Wrapf(err, "order by %d", i)
+			}
+			switch e.order {
+			case Asc:
+				p.Asc()
+			case Desc:
+				p.Desc()
+			default:
+				return errors.Errorf("unexpected order %v", e.order)
+			}
+		}
+	}
 
 	if q.limit > 0 {
 		p.Limit()
diff --git a/internal/chstorage/chsql/select_test.go b/internal/chstorage/chsql/select_test.go
index ec0ab0c7..659aa58a 100644
--- a/internal/chstorage/chsql/select_test.go
+++ b/internal/chstorage/chsql/select_test.go
@@ -76,10 +76,38 @@ func TestSelect(t *testing.T) {
 						Lt(Ident("duration"), Value[float32](3.14)),
 					),
 					Column("span_id", nil),
+				).
+					Limit(1)
+			},
+			false,
+		},
+		{
+			func() *SelectQuery {
+				return Select("spans",
+					Column("timestamp", nil),
+				).Order(
+					Ident("timestamp"),
+					Desc,
 				)
 			},
 			false,
 		},
+		{
+			func() *SelectQuery {
+				return Select("spans",
+					Column("timestamp", nil),
+				).Order(
+					Ident("timestamp"),
+					Asc,
+				).Order(
+					Ident("duration"),
+					Desc,
+				)
+			},
+			false,
+		},
+
+		// No columns.
 		{
 			func() *SelectQuery { return Select("logs") },
 			true,