Skip to content

Commit e117066

Browse files
johnstcnjohnstcn
committed
feat(scripts): add script to check schema between migrations
- migrations: allow passing in a custom migrate.FS - gen/dump: extract some functions to dbtestutil - scripts: write script to test migrations
1 parent d3f3ace commit e117066

File tree

4 files changed

+248
-102
lines changed

4 files changed

+248
-102
lines changed

‎coderd/database/dbtestutil/db.go

Lines changed: 79 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"os/exec"
1111
"path/filepath"
1212
"regexp"
13+
"strconv"
1314
"strings"
1415
"testing"
1516
"time"
@@ -184,20 +185,21 @@ func DumpOnFailure(t testing.TB, connectionURL string) {
184185
now := time.Now()
185186
timeSuffix := fmt.Sprintf("%d%d%d%d%d%d", now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
186187
outPath := filepath.Join(cwd, snakeCaseName+"."+timeSuffix+".test.sql")
187-
dump, err := pgDump(connectionURL)
188+
dump, err := PGDump(connectionURL)
188189
if err != nil {
189190
t.Errorf("dump on failure: failed to run pg_dump")
190191
return
191192
}
192-
if err := os.WriteFile(outPath, filterDump(dump), 0o600); err != nil {
193+
if err := os.WriteFile(outPath, normalizeDump(dump), 0o600); err != nil {
193194
t.Errorf("dump on failure: failed to write: %s", err.Error())
194195
return
195196
}
196197
t.Logf("Dumped database to %q due to failed test. I hope you find what you're looking for!", outPath)
197198
}
198199

199-
// pgDump runs pg_dump against dbURL and returns the output.
200-
func pgDump(dbURL string) ([]byte, error) {
200+
// PGDump runs pg_dump against dbURL and returns the output.
201+
// It is used by DumpOnFailure().
202+
func PGDump(dbURL string) ([]byte, error) {
201203
if _, err := exec.LookPath("pg_dump"); err != nil {
202204
return nil, xerrors.Errorf("could not find pg_dump in path: %w", err)
203205
}
@@ -230,16 +232,79 @@ func pgDump(dbURL string) ([]byte, error) {
230232
return stdout.Bytes(), nil
231233
}
232234

233-
// Unfortunately, some insert expressions span multiple lines.
234-
// The below may be over-permissive but better that than truncating data.
235-
var insertExpr = regexp.MustCompile(`(?s)\bINSERT[^;]+;`)
235+
const minimumPostgreSQLVersion = 13
236236

237-
func filterDump(dump []byte) []byte {
238-
var buf bytes.Buffer
239-
matches := insertExpr.FindAll(dump, -1)
240-
for _, m := range matches {
241-
_, _ = buf.Write(m)
242-
_, _ = buf.WriteRune('\n')
237+
// PGDumpSchemaOnly is for use by gen/dump only.
238+
// It runs pg_dump against dbURL and sets a consistent timezone and encoding.
239+
func PGDumpSchemaOnly(dbURL string) ([]byte, error) {
240+
hasPGDump := false
241+
if _, err := exec.LookPath("pg_dump"); err == nil {
242+
out, err := exec.Command("pg_dump", "--version").Output()
243+
if err == nil {
244+
// Parse output:
245+
// pg_dump (PostgreSQL) 14.5 (Ubuntu 14.5-0ubuntu0.22.04.1)
246+
parts := strings.Split(string(out), " ")
247+
if len(parts) > 2 {
248+
version, err := strconv.Atoi(strings.Split(parts[2], ".")[0])
249+
if err == nil && version >= minimumPostgreSQLVersion {
250+
hasPGDump = true
251+
}
252+
}
253+
}
243254
}
244-
return buf.Bytes()
255+
256+
cmdArgs := []string{
257+
"pg_dump",
258+
"--schema-only",
259+
dbURL,
260+
"--no-privileges",
261+
"--no-owner",
262+
"--no-privileges",
263+
"--no-publication",
264+
"--no-security-labels",
265+
"--no-subscriptions",
266+
"--no-tablespaces",
267+
268+
// We never want to manually generate
269+
// queries executing against this table.
270+
"--exclude-table=schema_migrations",
271+
}
272+
273+
if !hasPGDump {
274+
cmdArgs = append([]string{
275+
"docker",
276+
"run",
277+
"--rm",
278+
"--network=host",
279+
fmt.Sprintf("gcr.io/coder-dev-1/postgres:%d", minimumPostgreSQLVersion),
280+
}, cmdArgs...)
281+
}
282+
cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...) //#nosec
283+
cmd.Env = append(os.Environ(), []string{
284+
"PGTZ=UTC",
285+
"PGCLIENTENCODING=UTF8",
286+
}...)
287+
var output bytes.Buffer
288+
cmd.Stdout = &output
289+
cmd.Stderr = os.Stderr
290+
err := cmd.Run()
291+
if err != nil {
292+
return nil, err
293+
}
294+
return normalizeDump(output.Bytes()), nil
295+
}
296+
297+
func normalizeDump(schema []byte) []byte {
298+
// Remove all comments.
299+
schema = regexp.MustCompile(`(?im)^(--.*)$`).ReplaceAll(schema, []byte{})
300+
// Public is implicit in the schema.
301+
schema = regexp.MustCompile(`(?im)( |::|'|\()public\.`).ReplaceAll(schema, []byte(`$1`))
302+
// Remove database settings.
303+
schema = regexp.MustCompile(`(?im)^(SET.*;)`).ReplaceAll(schema, []byte(``))
304+
// Remove select statements
305+
schema = regexp.MustCompile(`(?im)^(SELECT.*;)`).ReplaceAll(schema, []byte(``))
306+
// Removes multiple newlines.
307+
schema = regexp.MustCompile(`(?im)\n{3,}`).ReplaceAll(schema, []byte("\n\n"))
308+
309+
return schema
245310
}

‎coderd/database/gen/dump/main.go

Lines changed: 4 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,16 @@
11
package main
22

33
import (
4-
"bytes"
54
"database/sql"
6-
"fmt"
75
"os"
8-
"os/exec"
96
"path/filepath"
107
"runtime"
11-
"strconv"
12-
"strings"
138

149
"github.com/coder/coder/v2/coderd/database/dbtestutil"
1510
"github.com/coder/coder/v2/coderd/database/migrations"
1611
)
1712

18-
const minimumPostgreSQLVersion = 13
13+
var preamble = []byte("-- Code generated by 'make coderd/database/generate'. DO NOT EDIT.")
1914

2015
func main() {
2116
connection, closeFn, err := dbtestutil.Open()
@@ -28,95 +23,23 @@ func main() {
2823
if err != nil {
2924
panic(err)
3025
}
26+
defer db.Close()
3127

3228
err = migrations.Up(db)
3329
if err != nil {
3430
panic(err)
3531
}
3632

37-
hasPGDump := false
38-
if _, err = exec.LookPath("pg_dump"); err == nil {
39-
out, err := exec.Command("pg_dump", "--version").Output()
40-
if err == nil {
41-
// Parse output:
42-
// pg_dump (PostgreSQL) 14.5 (Ubuntu 14.5-0ubuntu0.22.04.1)
43-
parts := strings.Split(string(out), " ")
44-
if len(parts) > 2 {
45-
version, err := strconv.Atoi(strings.Split(parts[2], ".")[0])
46-
if err == nil && version >= minimumPostgreSQLVersion {
47-
hasPGDump = true
48-
}
49-
}
50-
}
51-
}
52-
53-
cmdArgs := []string{
54-
"pg_dump",
55-
"--schema-only",
56-
connection,
57-
"--no-privileges",
58-
"--no-owner",
59-
60-
// We never want to manually generate
61-
// queries executing against this table.
62-
"--exclude-table=schema_migrations",
63-
}
64-
65-
if !hasPGDump {
66-
cmdArgs = append([]string{
67-
"docker",
68-
"run",
69-
"--rm",
70-
"--network=host",
71-
fmt.Sprintf("gcr.io/coder-dev-1/postgres:%d", minimumPostgreSQLVersion),
72-
}, cmdArgs...)
73-
}
74-
cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...) //#nosec
75-
cmd.Env = append(os.Environ(), []string{
76-
"PGTZ=UTC",
77-
"PGCLIENTENCODING=UTF8",
78-
}...)
79-
var output bytes.Buffer
80-
cmd.Stdout = &output
81-
cmd.Stderr = os.Stderr
82-
err = cmd.Run()
33+
dumpBytes, err := dbtestutil.PGDumpSchemaOnly(connection)
8334
if err != nil {
8435
panic(err)
8536
}
8637

87-
for _, sed := range []string{
88-
// Remove all comments.
89-
"/^--/d",
90-
// Public is implicit in the schema.
91-
"s/ public\\./ /g",
92-
"s/::public\\./::/g",
93-
"s/'public\\./'/g",
94-
"s/(public\\./(/g",
95-
// Remove database settings.
96-
"s/SET .* = .*;//g",
97-
// Remove select statements. These aren't useful
98-
// to a reader of the dump.
99-
"s/SELECT.*;//g",
100-
// Removes multiple newlines.
101-
"/^$/N;/^\\n$/D",
102-
} {
103-
cmd := exec.Command("sed", "-e", sed)
104-
cmd.Stdin = bytes.NewReader(output.Bytes())
105-
output = bytes.Buffer{}
106-
cmd.Stdout = &output
107-
cmd.Stderr = os.Stderr
108-
err = cmd.Run()
109-
if err != nil {
110-
panic(err)
111-
}
112-
}
113-
114-
dump := fmt.Sprintf("-- Code generated by 'make coderd/database/generate'. DO NOT EDIT.\n%s", output.Bytes())
11538
_, mainPath, _, ok := runtime.Caller(0)
11639
if !ok {
11740
panic("couldn't get caller path")
11841
}
119-
err = os.WriteFile(filepath.Join(mainPath, "..", "..", "..", "dump.sql"), []byte(dump), 0o600)
42+
err = os.WriteFile(filepath.Join(mainPath, "..", "..", "..", "dump.sql"), append(preamble, dumpBytes...), 0o600)
12043
if err != nil {
12144
panic(err)
12245
}

‎coderd/database/migrations/migrate.go

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@ import (
1717
//go:embed *.sql
1818
var migrations embed.FS
1919

20-
func setup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
20+
func setup(db *sql.DB, migs fs.FS) (source.Driver, *migrate.Migrate, error) {
21+
if migs == nil {
22+
migs = migrations
23+
}
2124
ctx := context.Background()
22-
sourceDriver, err := iofs.New(migrations, ".")
25+
sourceDriver, err := iofs.New(migs, ".")
2326
if err != nil {
2427
return nil, nil, xerrors.Errorf("create iofs: %w", err)
2528
}
@@ -47,8 +50,13 @@ func setup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
4750
}
4851

4952
// Up runs SQL migrations to ensure the database schema is up-to-date.
50-
func Up(db *sql.DB) (retErr error) {
51-
_, m, err := setup(db)
53+
func Up(db *sql.DB) error {
54+
return UpWithFS(db, migrations)
55+
}
56+
57+
// UpWithFS runs SQL migrations in the given fs.
58+
func UpWithFS(db *sql.DB, migs fs.FS) (retErr error) {
59+
_, m, err := setup(db, migs)
5260
if err != nil {
5361
return xerrors.Errorf("migrate setup: %w", err)
5462
}
@@ -79,7 +87,7 @@ func Up(db *sql.DB) (retErr error) {
7987

8088
// Down runs all down SQL migrations.
8189
func Down(db *sql.DB) error {
82-
_, m, err := setup(db)
90+
_, m, err := setup(db, migrations)
8391
if err != nil {
8492
return xerrors.Errorf("migrate setup: %w", err)
8593
}
@@ -101,7 +109,7 @@ func Down(db *sql.DB) error {
101109
// applied, without making any changes to the database. If not, returns a
102110
// non-nil error.
103111
func EnsureClean(db *sql.DB) error {
104-
sourceDriver, m, err := setup(db)
112+
sourceDriver, m, err := setup(db, migrations)
105113
if err != nil {
106114
return xerrors.Errorf("migrate setup: %w", err)
107115
}
@@ -167,7 +175,7 @@ func CheckLatestVersion(sourceDriver source.Driver, currentVersion uint) error {
167175
// Stepper cannot be closed pre-emptively, it must be run to completion
168176
// (or until an error is encountered).
169177
func Stepper(db *sql.DB) (next func() (version uint, more bool, err error), err error) {
170-
_, m, err := setup(db)
178+
_, m, err := setup(db, migrations)
171179
if err != nil {
172180
return nil, xerrors.Errorf("migrate setup: %w", err)
173181
}

0 commit comments

Comments
 (0)