From bb8f6a0f2c5cbbb2f194f2fd9f5505cde1fe5776 Mon Sep 17 00:00:00 2001
From: Kirill Zaitsev <kirik910@gmail.com>
Date: Sat, 19 Aug 2023 12:07:15 +0300
Subject: [PATCH] feat: union expr

---
 ch/query_select.go                 | 47 +++++++++++++++++++++++++++++-
 ch/query_test.go                   | 10 +++++++
 ch/testdata/snapshots/TestQuery-17 |  1 +
 ch/testdata/snapshots/TestQuery-18 |  1 +
 4 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 ch/testdata/snapshots/TestQuery-17
 create mode 100644 ch/testdata/snapshots/TestQuery-18

diff --git a/ch/query_select.go b/ch/query_select.go
index 92c66a2..fa7dbc7 100644
--- a/ch/query_select.go
+++ b/ch/query_select.go
@@ -12,6 +12,11 @@ import (
 	"github.com/uptrace/go-clickhouse/ch/internal"
 )
 
+type union struct {
+	expr  string
+	query *SelectQuery
+}
+
 type SelectQuery struct {
 	baseQuery
 
@@ -26,6 +31,8 @@ type SelectQuery struct {
 	limit      int
 	offset     int
 	final      bool
+
+	union []union
 }
 
 var _ Query = (*SelectQuery)(nil)
@@ -150,6 +157,24 @@ func (q *SelectQuery) ExcludeColumn(columns ...string) *SelectQuery {
 	return q
 }
 
+///------------------------------------------------------------------------------
+
+func (q *SelectQuery) Union(other *SelectQuery) *SelectQuery {
+	return q.addUnion(" UNION ", other)
+}
+
+func (q *SelectQuery) UnionAll(other *SelectQuery) *SelectQuery {
+	return q.addUnion(" UNION ALL ", other)
+}
+
+func (q *SelectQuery) addUnion(expr string, other *SelectQuery) *SelectQuery {
+	q.union = append(q.union, union{
+		expr:  expr,
+		query: other,
+	})
+	return q
+}
+
 //------------------------------------------------------------------------------
 
 func (q *SelectQuery) Join(join string, args ...any) *SelectQuery {
@@ -329,6 +354,10 @@ func (q *SelectQuery) appendQuery(
 		b = append(b, `WITH "_count_wrapper" AS (`...)
 	}
 
+	if len(q.union) > 0 {
+		b = append(b, '(')
+	}
+
 	if len(q.with) > 0 {
 		b, err = q.appendWith(fmter, b)
 		if err != nil {
@@ -429,6 +458,20 @@ func (q *SelectQuery) appendQuery(
 		}
 	}
 
+	if len(q.union) > 0 {
+		b = append(b, ')')
+
+		for _, u := range q.union {
+			b = append(b, u.expr...)
+			b = append(b, '(')
+			b, err = u.query.AppendQuery(fmter, b)
+			if err != nil {
+				return nil, err
+			}
+			b = append(b, ')')
+		}
+	}
+
 	if !count {
 		if len(q.order) > 0 {
 			b = append(b, " ORDER BY "...)
@@ -451,7 +494,9 @@ func (q *SelectQuery) appendQuery(
 			b = append(b, " OFFSET "...)
 			b = strconv.AppendInt(b, int64(q.offset), 10)
 		}
-	} else if cteCount {
+	}
+
+	if cteCount {
 		b = append(b, `) SELECT `...)
 		b = append(b, "count()"...)
 		b = append(b, ` FROM "_count_wrapper"`...)
diff --git a/ch/query_test.go b/ch/query_test.go
index b3cbb62..772ffcd 100644
--- a/ch/query_test.go
+++ b/ch/query_test.go
@@ -115,6 +115,16 @@ func TestQuery(t *testing.T) {
 				Group("id").
 				OrderExpr("id")
 		},
+		func(db *ch.DB) chschema.QueryAppender {
+			q1 := db.NewSelect().Model(new(Model)).Where("1")
+			q2 := db.NewSelect().Model(new(Model))
+			return q1.Union(q2)
+		},
+		func(db *ch.DB) chschema.QueryAppender {
+			q1 := db.NewSelect().Model(new(Model)).Where("1")
+			q2 := db.NewSelect().Model(new(Model))
+			return q1.UnionAll(q2)
+		},
 	}
 
 	db := chDB()
diff --git a/ch/testdata/snapshots/TestQuery-17 b/ch/testdata/snapshots/TestQuery-17
new file mode 100644
index 0000000..3ca11e8
--- /dev/null
+++ b/ch/testdata/snapshots/TestQuery-17
@@ -0,0 +1 @@
+(SELECT "model"."id", "model"."string", "model"."bytes" FROM "models" AS "model" WHERE (1)) UNION (SELECT "model"."id", "model"."string", "model"."bytes" FROM "models" AS "model")
diff --git a/ch/testdata/snapshots/TestQuery-18 b/ch/testdata/snapshots/TestQuery-18
new file mode 100644
index 0000000..0a90f07
--- /dev/null
+++ b/ch/testdata/snapshots/TestQuery-18
@@ -0,0 +1 @@
+(SELECT "model"."id", "model"."string", "model"."bytes" FROM "models" AS "model" WHERE (1)) UNION ALL (SELECT "model"."id", "model"."string", "model"."bytes" FROM "models" AS "model")