Skip to content

Commit 02cb84e

Browse files
yhabteabyhabteab
committed
Introduce QueryBuilder type
1 parent f0e227c commit 02cb84e

File tree

2 files changed

+397
-0
lines changed

2 files changed

+397
-0
lines changed

‎database/query_builder.go‎

Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
package database
2+
3+
import (
4+
"fmt"
5+
"github.com/icinga/icinga-go-library/driver"
6+
"go.uber.org/zap"
7+
"golang.org/x/exp/slices"
8+
"reflect"
9+
"sort"
10+
"strings"
11+
)
12+
13+
// QueryBuilder is an addon for the DB type that takes care of all the database statement building shenanigans.
14+
// Note: This type is designed primarily for one-off use (monouso) and subsequent disposal and should only be
15+
// used to generate a single database query type.
16+
type QueryBuilder struct {
17+
subject any
18+
columns []string
19+
excludedColumns []string
20+
21+
// Indicates whether the generated columns should be sorted in ascending order before generating the
22+
// actual statements. This is intended for unit tests only and shouldn't be necessary for production code.
23+
sort bool
24+
}
25+
26+
// NewQB returns a fully initialized *QueryBuilder instance for the given subject/struct.
27+
func NewQB(subject any) *QueryBuilder {
28+
return &QueryBuilder{subject: subject}
29+
}
30+
31+
// SetColumns sets the DB columns to be used when building the statements.
32+
// When you do not want the columns to be extracted dynamically, you can use this method to specify them manually.
33+
func (qb *QueryBuilder) SetColumns(columns ...string) {
34+
qb.columns = columns
35+
}
36+
37+
// SetExcludedColumns excludes the given columns from all the database statements.
38+
func (qb *QueryBuilder) SetExcludedColumns(columns ...string) {
39+
qb.excludedColumns = columns
40+
}
41+
42+
// Delete returns a DELETE statement for the query builders subject filtered by ID.
43+
func (qb *QueryBuilder) Delete() string {
44+
return qb.DeleteBy("id")
45+
}
46+
47+
// DeleteBy returns a DELETE statement for the query builders subject filtered by the given column.
48+
func (qb *QueryBuilder) DeleteBy(column string) string {
49+
return fmt.Sprintf(`DELETE FROM "%s" WHERE "%s" IN (?)`, TableName(qb.subject), column)
50+
}
51+
52+
// Insert returns an INSERT INTO statement for the query builders subject.
53+
func (qb *QueryBuilder) Insert(db *DB) (string, int) {
54+
columns := qb.BuildColumns(db)
55+
56+
return fmt.Sprintf(
57+
`INSERT INTO "%s" ("%s") VALUES (%s)`,
58+
TableName(qb.subject),
59+
strings.Join(columns, `", "`),
60+
fmt.Sprintf(":%s", strings.Join(columns, ", :")),
61+
), len(columns)
62+
}
63+
64+
// InsertIgnore returns an INSERT statement for the query builders subject for
65+
// which the database ignores rows that have already been inserted.
66+
func (qb *QueryBuilder) InsertIgnore(db *DB) (string, int) {
67+
columns := qb.BuildColumns(db)
68+
table := TableName(qb.subject)
69+
70+
var clause string
71+
switch db.DriverName() {
72+
case driver.MySQL:
73+
// MySQL treats UPDATE id = id as a no-op.
74+
clause = fmt.Sprintf(`ON DUPLICATE KEY UPDATE "%s" = "%s"`, columns[0], columns[0])
75+
case driver.PostgreSQL:
76+
var constraint string
77+
if constrainter, ok := qb.subject.(PgsqlOnConflictConstrainter); ok {
78+
constraint = constrainter.PgsqlOnConflictConstraint()
79+
} else {
80+
constraint = "pk_" + table
81+
}
82+
83+
clause = fmt.Sprintf("ON CONFLICT ON CONSTRAINT %s DO NOTHING", constraint)
84+
default:
85+
db.logger.Fatalw("Driver unsupported", zap.String("driver", db.DriverName()))
86+
}
87+
88+
return fmt.Sprintf(
89+
`INSERT INTO "%s" ("%s") VALUES (%s) %s`,
90+
table,
91+
strings.Join(columns, `", "`),
92+
fmt.Sprintf(":%s", strings.Join(columns, ", :")),
93+
clause,
94+
), len(columns)
95+
}
96+
97+
// Select returns a SELECT statement from the query builders subject and the already set columns.
98+
// If no columns are set, they will be extracted from the query builders subject.
99+
// When the query builders subject is of type Scoper, a WHERE clause is appended to the statement.
100+
func (qb *QueryBuilder) Select(db *DB) string {
101+
var scoper Scoper
102+
if sc, ok := qb.subject.(Scoper); ok {
103+
scoper = sc
104+
}
105+
106+
return qb.SelectScoped(db, scoper)
107+
}
108+
109+
// SelectScoped returns a SELECT statement from the query builders subject and the already set columns filtered
110+
// by the given scoper/column. When no columns are set, they will be extracted from the query builders subject.
111+
// The argument scoper must either be of type Scoper, string or nil to get SELECT statements without a WHERE clause.
112+
func (qb *QueryBuilder) SelectScoped(db *DB, scoper any) string {
113+
query := fmt.Sprintf(`SELECT "%s" FROM "%s"`, strings.Join(qb.BuildColumns(db), `", "`), TableName(qb.subject))
114+
where, placeholders := qb.Where(db, scoper)
115+
if placeholders > 0 {
116+
query += ` WHERE ` + where
117+
}
118+
119+
return query
120+
}
121+
122+
// Update returns an UPDATE statement for the query builders subject filter by ID column.
123+
func (qb *QueryBuilder) Update(db *DB) (string, int) {
124+
return qb.UpdateScoped(db, "id")
125+
}
126+
127+
// UpdateScoped returns an UPDATE statement for the query builders subject filtered by the given column/scoper.
128+
// The argument scoper must either be of type Scoper, string or nil to get UPDATE statements without a WHERE clause.
129+
func (qb *QueryBuilder) UpdateScoped(db *DB, scoper any) (string, int) {
130+
columns := qb.BuildColumns(db)
131+
set := make([]string, 0, len(columns))
132+
133+
for _, col := range columns {
134+
set = append(set, fmt.Sprintf(`"%s" = :%s`, col, col))
135+
}
136+
137+
placeholders := len(columns)
138+
query := `UPDATE "%s" SET %s`
139+
if where, count := qb.Where(db, scoper); count > 0 {
140+
placeholders += count
141+
query += ` WHERE ` + where
142+
}
143+
144+
return fmt.Sprintf(query, TableName(qb.subject), strings.Join(set, ", ")), placeholders
145+
}
146+
147+
// Upsert returns an upsert statement for the query builders subject.
148+
func (qb *QueryBuilder) Upsert(db *DB) (string, int) {
149+
var updateColumns []string
150+
if upserter, ok := qb.subject.(Upserter); ok {
151+
updateColumns = db.BuildColumns(upserter.Upsert())
152+
} else {
153+
updateColumns = qb.BuildColumns(db)
154+
}
155+
156+
return qb.UpsertColumns(db, updateColumns...)
157+
}
158+
159+
// UpsertColumns returns an upsert statement for the query builders subject and the specified update columns.
160+
func (qb *QueryBuilder) UpsertColumns(db *DB, updateColumns ...string) (string, int) {
161+
insertColumns := qb.BuildColumns(db)
162+
table := TableName(qb.subject)
163+
164+
var clause, setFormat string
165+
switch db.DriverName() {
166+
case driver.MySQL:
167+
clause = "ON DUPLICATE KEY UPDATE"
168+
setFormat = `"%[1]s" = VALUES("%[1]s")`
169+
case driver.PostgreSQL:
170+
var constraint string
171+
if constrainter, ok := qb.subject.(PgsqlOnConflictConstrainter); ok {
172+
constraint = constrainter.PgsqlOnConflictConstraint()
173+
} else {
174+
constraint = "pk_" + table
175+
}
176+
177+
clause = fmt.Sprintf("ON CONFLICT ON CONSTRAINT %s DO UPDATE SET", constraint)
178+
setFormat = `"%[1]s" = EXCLUDED."%[1]s"`
179+
default:
180+
db.logger.Fatalw("Driver unsupported", zap.String("driver", db.DriverName()))
181+
}
182+
183+
set := make([]string, 0, len(updateColumns))
184+
for _, col := range updateColumns {
185+
set = append(set, fmt.Sprintf(setFormat, col))
186+
}
187+
188+
return fmt.Sprintf(
189+
`INSERT INTO "%s" ("%s") VALUES (%s) %s %s`,
190+
table,
191+
strings.Join(insertColumns, `", "`),
192+
fmt.Sprintf(":%s", strings.Join(insertColumns, ", :")),
193+
clause,
194+
strings.Join(set, ", "),
195+
), len(insertColumns)
196+
}
197+
198+
// Where returns a WHERE clause with named placeholder conditions built from the
199+
// specified scoper/column combined with the AND operator.
200+
func (qb *QueryBuilder) Where(db *DB, subject any) (string, int) {
201+
var columns []string
202+
t := reflect.TypeOf(subject)
203+
if t.Kind() == reflect.String {
204+
columns = []string{subject.(string)}
205+
} else if t.Kind() == reflect.Struct || t.Kind() == reflect.Pointer {
206+
if scoper, ok := subject.(Scoper); ok {
207+
return qb.Where(db, scoper.Scope())
208+
}
209+
210+
columns = db.BuildColumns(subject)
211+
}
212+
213+
where := make([]string, 0, len(columns))
214+
for _, col := range columns {
215+
where = append(where, fmt.Sprintf(`"%s" = :%s`, col, col))
216+
}
217+
218+
return strings.Join(where, ` AND `), len(columns)
219+
}
220+
221+
// BuildColumns returns all the Query Builder columns (if specified), otherwise they are
222+
// determined dynamically using its subject. Additionally, it checks whether columns need
223+
// to be excluded and proceeds accordingly.
224+
func (qb *QueryBuilder) BuildColumns(db *DB) []string {
225+
var columns []string
226+
if len(qb.columns) > 0 {
227+
columns = qb.columns
228+
} else {
229+
columns = db.BuildColumns(qb.subject)
230+
}
231+
232+
if len(qb.excludedColumns) > 0 {
233+
columns = slices.DeleteFunc(append([]string(nil), columns...), func(column string) bool {
234+
for _, exclude := range qb.excludedColumns {
235+
if exclude == column {
236+
return true
237+
}
238+
}
239+
240+
return false
241+
})
242+
}
243+
244+
if qb.sort {
245+
// The order in which the columns appear is not guaranteed as we extract the columns dynamically
246+
// from the struct. So, we've to sort them here to be able to test the generated statements.
247+
sort.SliceStable(columns, func(a, b int) bool {
248+
return columns[a] < columns[b]
249+
})
250+
}
251+
252+
return columns
253+
}

‎database/query_builder_test.go‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
package database
2+
3+
import (
4+
"github.com/icinga/icinga-go-library/driver"
5+
"github.com/stretchr/testify/require"
6+
"testing"
7+
)
8+
9+
func TestQueryBuilder(t *testing.T) {
10+
t.Parallel()
11+
12+
conf := Config{
13+
Type: "mysql",
14+
Host: "localhost",
15+
Port: 3306,
16+
Database: "igl-test",
17+
User: "igl-test",
18+
Password: "igl-test",
19+
}
20+
21+
driver.Register(nil)
22+
db, err := NewDbFromConfig(&conf, nil)
23+
require.NoError(t, err)
24+
25+
t.Run("SetColumns", func(t *testing.T) {
26+
qb := NewQB("test")
27+
qb.SetColumns("a", "b")
28+
require.Equal(t, []string{"a", "b"}, qb.columns)
29+
})
30+
31+
t.Run("ExcludeColumns", func(t *testing.T) {
32+
qb := NewQB(&test{})
33+
qb.SetExcludedColumns("a", "b")
34+
require.Equal(t, []string{"a", "b"}, qb.excludedColumns)
35+
})
36+
37+
t.Run("DeleteStatements", func(t *testing.T) {
38+
qb := NewQB(&test{})
39+
require.Equal(t, `DELETE FROM "test" WHERE "id" IN (?)`, qb.Delete())
40+
require.Equal(t, `DELETE FROM "test" WHERE "foo" IN (?)`, qb.DeleteBy("foo"))
41+
})
42+
43+
t.Run("InsertStatements", func(t *testing.T) {
44+
t.Parallel()
45+
46+
qb := NewQB(&test{})
47+
qb.sort = true
48+
qb.SetExcludedColumns("random")
49+
50+
stmt, columns := qb.Insert(db)
51+
require.Equal(t, 2, columns)
52+
require.Equal(t, `INSERT INTO "test" ("name", "value") VALUES (:name, :value)`, stmt)
53+
54+
qb.SetExcludedColumns("a", "b")
55+
qb.SetColumns("a", "b", "c", "d")
56+
57+
stmt, columns = qb.Insert(db)
58+
require.Equal(t, 2, columns)
59+
require.Equal(t, `INSERT INTO "test" ("c", "d") VALUES (:c, :d)`, stmt)
60+
61+
stmt, columns = qb.InsertIgnore(db)
62+
require.Equal(t, 2, columns)
63+
require.Equal(t, `INSERT INTO "test" ("c", "d") VALUES (:c, :d) ON DUPLICATE KEY UPDATE "c" = "c"`, stmt)
64+
})
65+
66+
t.Run("SelectStatements", func(t *testing.T) {
67+
t.Parallel()
68+
69+
qb := NewQB(&test{})
70+
qb.sort = true
71+
72+
stmt := qb.Select(db)
73+
expected := `SELECT "name", "random", "value" FROM "test" WHERE "scoper_id" = :scoper_id`
74+
require.Equal(t, expected, stmt)
75+
76+
qb.SetColumns("name", "random", "value")
77+
78+
stmt = qb.SelectScoped(db, "name")
79+
require.Equal(t, `SELECT "name", "random", "value" FROM "test" WHERE "name" = :name`, stmt)
80+
})
81+
82+
t.Run("UpdateStatements", func(t *testing.T) {
83+
t.Parallel()
84+
85+
qb := NewQB(&test{})
86+
qb.sort = true
87+
qb.SetExcludedColumns("random")
88+
89+
stmt, placeholders := qb.Update(db)
90+
require.Equal(t, 3, placeholders)
91+
92+
expected := `UPDATE "test" SET "name" = :name, "value" = :value WHERE "id" = :id`
93+
require.Equal(t, expected, stmt)
94+
95+
stmt, placeholders = qb.UpdateScoped(db, (&test{}).Scope())
96+
require.Equal(t, 3, placeholders)
97+
require.Equal(t, `UPDATE "test" SET "name" = :name, "value" = :value WHERE "scoper_id" = :scoper_id`, stmt)
98+
99+
qb.SetExcludedColumns("a", "b")
100+
qb.SetColumns("a", "b", "c", "d")
101+
102+
stmt, placeholders = qb.UpdateScoped(db, "c")
103+
require.Equal(t, 3, placeholders)
104+
require.Equal(t, 3, placeholders)
105+
require.Equal(t, `UPDATE "test" SET "c" = :c, "d" = :d WHERE "c" = :c`, stmt)
106+
})
107+
108+
t.Run("UpsertStatements", func(t *testing.T) {
109+
t.Parallel()
110+
111+
qb := NewQB(&test{})
112+
qb.sort = true
113+
qb.SetExcludedColumns("random")
114+
115+
stmt, columns := qb.Upsert(db)
116+
require.Equal(t, 2, columns)
117+
require.Equal(t, `INSERT INTO "test" ("name", "value") VALUES (:name, :value) ON DUPLICATE KEY UPDATE "name" = VALUES("name"), "value" = VALUES("value")`, stmt)
118+
119+
qb.SetExcludedColumns("a", "b")
120+
qb.SetColumns("a", "b", "c", "d")
121+
122+
stmt, columns = qb.Upsert(db)
123+
require.Equal(t, 2, columns)
124+
require.Equal(t, `INSERT INTO "test" ("c", "d") VALUES (:c, :d) ON DUPLICATE KEY UPDATE "c" = VALUES("c"), "d" = VALUES("d")`, stmt)
125+
126+
qb.SetExcludedColumns("a")
127+
128+
stmt, columns = qb.UpsertColumns(db, "b", "c")
129+
require.Equal(t, 3, columns)
130+
require.Equal(t, `INSERT INTO "test" ("b", "c", "d") VALUES (:b, :c, :d) ON DUPLICATE KEY UPDATE "b" = VALUES("b"), "c" = VALUES("c")`, stmt)
131+
})
132+
}
133+
134+
type test struct {
135+
Name string
136+
Value string
137+
Random string
138+
}
139+
140+
func (t *test) Scope() any {
141+
return struct {
142+
ScoperID string
143+
}{}
144+
}

0 commit comments

Comments
 (0)