From 4955695c363a15d8be926d604dc48af0b7940554 Mon Sep 17 00:00:00 2001
From: Vlad Glushchuk <vgluschuk@gmail.com>
Date: Wed, 11 Sep 2019 17:20:10 +0200
Subject: [PATCH] Add Union method

---
 README.md       | 17 ++++++++++++++++-
 example_test.go | 20 ++++++++++++++++++++
 stmt.go         | 36 ++++++++++++++++++++++++++++++++++++
 stmt_test.go    | 11 +++++++++++
 4 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index beff640..7d1a2b6 100644
--- a/README.md
+++ b/README.md
@@ -229,7 +229,7 @@ Use `SubQuery` method to add a sub query to a statement:
 	q.Close()
 ```
 
-Not that if a subquery uses no arguments, it's more effective to add it as SQL fragment:
+Note that if a subquery uses no arguments, it's more effective to add it as SQL fragment:
 
 ```go
 	q := sqlf.From("orders o").
@@ -241,6 +241,21 @@ Not that if a subquery uses no arguments, it's more effective to add it as SQL f
     q.Close()
 ```
 
+#### Unions
+
+Use `Union` method to combine results of two queries:
+
+```go
+	q := sqlf.From("tasks").
+		Select("id, status").
+		Where("status = ?", "new").
+		Union(false, sqlf.PostgreSQL.From("tasks").
+			Select("id, status").
+            Where("status = ?", "wip"))
+    // ...
+	q.Close()
+```
+
 ### INSERT
 
 `sqlf` provides a `Set` method to be used both for UPDATE and INSERT statements:
diff --git a/example_test.go b/example_test.go
index 8ae8a91..70c2c53 100644
--- a/example_test.go
+++ b/example_test.go
@@ -267,3 +267,23 @@ func ExampleStmt_In() {
 	// SELECT id, status FROM tasks WHERE status IN (?,?,?)
 	// [new pending wip]
 }
+
+func ExampleStmt_Union() {
+	q := sqlf.From("tasks").
+		Select("id, status").
+		Where("status = ?", "new").
+		Union(true, sqlf.From("tasks").
+			Select("id, status").
+			Where("status = ?", "pending")).
+		Union(true, sqlf.From("tasks").
+			Select("id, status").
+			Where("status = ?", "wip")).
+		OrderBy("id")
+	fmt.Println(q.String())
+	fmt.Println(q.Args())
+	q.Close()
+
+	// Output:
+	// SELECT id, status FROM tasks WHERE status = ? UNION ALL SELECT id, status FROM tasks WHERE status = ? UNION ALL SELECT id, status FROM tasks WHERE status = ? ORDER BY id
+	// [new pending wip]
+}
diff --git a/stmt.go b/stmt.go
index 19f750c..0992717 100644
--- a/stmt.go
+++ b/stmt.go
@@ -477,6 +477,41 @@ func (q *Stmt) SubQuery(prefix, suffix string, query *Stmt) *Stmt {
 	return q
 }
 
+/*
+Union adds a UNION clause to the statement.
+
+all argument controls if UNION ALL or UNION clause
+is to be constructed. Use UNION ALL if possible to
+get faster queries.
+*/
+func (q *Stmt) Union(all bool, query *Stmt) *Stmt {
+	p := posUnion
+	if len(q.chunks) > 0 {
+		last := (&q.chunks[len(q.chunks)-1]).pos
+		if last >= p {
+			p = last + 1
+		}
+	}
+	var index int
+	if all {
+		index = q.addChunk(p, "UNION ALL ", "", query.args, "")
+	} else {
+		index = q.addChunk(p, "UNION ", "", query.args, "")
+	}
+	chunk := &q.chunks[index]
+	// Make sure subquery is not dialect-specific.
+	if query.dialect != NoDialect {
+		query.dialect = NoDialect
+		query.Invalidate()
+	}
+	q.buf.WriteString(query.String())
+	chunk.bufHigh = q.buf.Len()
+	// Close the subquery
+	query.Close()
+
+	return q
+}
+
 /*
 Clause appends a raw SQL fragment to the statement.
 
@@ -725,6 +760,7 @@ const (
 	posWhere
 	posGroupBy
 	posHaving
+	posUnion
 	posOrderBy
 	posLimit
 	posOffset
diff --git a/stmt_test.go b/stmt_test.go
index d2058e3..3cb250b 100644
--- a/stmt_test.go
+++ b/stmt_test.go
@@ -233,3 +233,14 @@ func TestFullJoin(t *testing.T) {
 	defer q.Close()
 	assert.Equal(t, "SELECT id FROM orders o FULL JOIN users u ON (u.id = o.user_id)", q.String())
 }
+
+func TestUnion(t *testing.T) {
+	q := sqlf.From("tasks").
+		Select("id, status").
+		Where("status = ?", "new").
+		Union(false, sqlf.PostgreSQL.From("tasks").
+			Select("id, status").
+			Where("status = ?", "wip"))
+	defer q.Close()
+	assert.Equal(t, "SELECT id, status FROM tasks WHERE status = ? UNION SELECT id, status FROM tasks WHERE status = ?", q.String())
+}