Skip to content

Commit 8c20d9d

Browse files
coadlercoadler
committed
feat: run all migrations in a transaction
1 parent 5b2f436 commit 8c20d9d

10 files changed

+231
-18
lines changed

‎coderd/database/dump.sql

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000030_template_version_created_by.up.sql

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
BEGIN;
21

32
ALTER TABLE ONLY template_versions ADD COLUMN IF NOT EXISTS created_by uuid REFERENCES users (id) ON DELETE RESTRICT;
43

‎coderd/database/migrations/000035_linked_user_id.down.sql

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22
-- the oauth_access_token, oauth_refresh_token, and oauth_expiry
33
-- columns of api_key rows with the values from the dropped user_links
44
-- table.
5-
BEGIN;
6-
75
DROP TABLE IF EXISTS user_links;
86

97
ALTER TABLE
@@ -19,5 +17,3 @@ ALTER TABLE
1917
ADD COLUMN oauth_expiry timestamp with time zone DEFAULT '0001-01-01 00:00:00+00'::timestamp with time zone NOT NULL;
2018

2119
ALTER TABLE users DROP COLUMN login_type;
22-
23-
COMMIT;

‎coderd/database/migrations/000057_api_key_token.up.sql

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,22 @@
1-
ALTER TYPE login_type ADD VALUE IF NOT EXISTS 'token';
1+
-- ALTER TYPE login_type ADD VALUE IF NOT EXISTS 'token';
2+
3+
CREATE TYPE new_logintype AS ENUM (
4+
'password',
5+
'github',
6+
'oidc',
7+
'token'
8+
);
9+
10+
ALTER TABLE users
11+
ALTER COLUMN login_type DROP DEFAULT,
12+
ALTER COLUMN login_type TYPE new_logintype USING (login_type::text::new_logintype),
13+
ALTER COLUMN login_type SET DEFAULT 'password'::new_logintype;
14+
15+
ALTER TABLE user_links
16+
ALTER COLUMN login_type TYPE new_logintype USING (login_type::text::new_logintype);
17+
18+
ALTER TABLE api_keys
19+
ALTER COLUMN login_type TYPE new_logintype USING (login_type::text::new_logintype);
20+
21+
DROP TYPE login_type;
22+
ALTER TYPE new_logintype RENAME TO login_type;

‎coderd/database/migrations/000144_user_status_dormant.down.sql

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
-- It's not possible to drop enum values from enum types, so the UP has "IF NOT EXISTS"
22

3-
UPDATE users SET status = 'active'::user_status WHERE status = 'dormant'::user_status;
3+
UPDATE users SET status = 'active'::user_status WHERE status::text = 'dormant';

‎coderd/database/migrations/000144_user_status_dormant.up.sql

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,14 @@
1-
ALTER TYPE user_status ADD VALUE IF NOT EXISTS 'dormant';
2-
COMMENT ON TYPE user_status IS 'Defines the user status: active, dormant, or suspended.';
1+
CREATE TYPE new_user_status AS ENUM (
2+
'active',
3+
'suspended',
4+
'dormant'
5+
);
6+
COMMENT ON TYPE new_user_status IS 'Defines the users status: active, dormant, or suspended.';
7+
8+
ALTER TABLE users
9+
ALTER COLUMN status DROP DEFAULT,
10+
ALTER COLUMN status TYPE new_user_status USING (status::text::new_user_status),
11+
ALTER COLUMN status SET DEFAULT 'active'::new_user_status;
12+
13+
DROP TYPE user_status;
14+
ALTER TYPE new_user_status RENAME TO user_status;

‎coderd/database/migrations/migrate.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,9 @@ func setup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
2727

2828
// migration_cursor is a v1 migration table. If this exists, we're on v1.
2929
// Do no run v2 migrations on a v1 database!
30-
row := db.QueryRowContext(ctx, "SELECT * FROM migration_cursor;")
31-
if row.Err() == nil {
30+
row := db.QueryRowContext(ctx, "SELECT 1 FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = 'migration_cursor';")
31+
var v1Exists int
32+
if row.Scan(&v1Exists) == nil {
3233
return nil, nil, xerrors.Errorf("currently connected to a Coder v1 database, aborting database setup")
3334
}
3435

@@ -55,7 +56,7 @@ func setup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
5556

5657
// Up runs SQL migrations to ensure the database schema is up-to-date.
5758
func Up(db *sql.DB) (retErr error) {
58-
_, m, err := setup(db)
59+
_, m, err := betterSetup(db)
5960
if err != nil {
6061
return xerrors.Errorf("migrate setup: %w", err)
6162
}
@@ -86,7 +87,7 @@ func Up(db *sql.DB) (retErr error) {
8687

8788
// Down runs all down SQL migrations.
8889
func Down(db *sql.DB) error {
89-
_, m, err := setup(db)
90+
_, m, err := betterSetup(db)
9091
if err != nil {
9192
return xerrors.Errorf("migrate setup: %w", err)
9293
}
@@ -108,7 +109,7 @@ func Down(db *sql.DB) error {
108109
// applied, without making any changes to the database. If not, returns a
109110
// non-nil error.
110111
func EnsureClean(db *sql.DB) error {
111-
sourceDriver, m, err := setup(db)
112+
sourceDriver, m, err := betterSetup(db)
112113
if err != nil {
113114
return xerrors.Errorf("migrate setup: %w", err)
114115
}
@@ -174,7 +175,7 @@ func CheckLatestVersion(sourceDriver source.Driver, currentVersion uint) error {
174175
// Stepper cannot be closed pre-emptively, it must be run to completion
175176
// (or until an error is encountered).
176177
func Stepper(db *sql.DB) (next func() (version uint, more bool, err error), err error) {
177-
_, m, err := setup(db)
178+
_, m, err := betterSetup(db)
178179
if err != nil {
179180
return nil, xerrors.Errorf("migrate setup: %w", err)
180181
}

‎coderd/database/migrations/txnmigrator.go

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
package migrations
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"database/sql"
7+
"fmt"
8+
"io"
9+
"strings"
10+
11+
"github.com/golang-migrate/migrate/v4"
12+
"github.com/golang-migrate/migrate/v4/database"
13+
"github.com/golang-migrate/migrate/v4/source"
14+
"github.com/golang-migrate/migrate/v4/source/iofs"
15+
"github.com/lib/pq"
16+
"golang.org/x/xerrors"
17+
)
18+
19+
const (
20+
lockID = int64(1037453835920848937)
21+
migrationsTableName = "schema_migrations"
22+
)
23+
24+
func betterSetup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
25+
ctx := context.Background()
26+
sourceDriver, err := iofs.New(migrations, ".")
27+
if err != nil {
28+
return nil, nil, xerrors.Errorf("create iofs: %w", err)
29+
}
30+
31+
// migration_cursor is a v1 migration table. If this exists, we're on v1.
32+
// Do no run v2 migrations on a v1 database!
33+
row := db.QueryRowContext(ctx, "SELECT 1 FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = 'migration_cursor';")
34+
var v1Exists int
35+
if row.Scan(&v1Exists) == nil {
36+
return nil, nil, xerrors.New("currently connected to a Coder v1 database, aborting database setup")
37+
}
38+
39+
dbDriver := &pgTxnDriver{ctx: context.Background(), db: db}
40+
err = dbDriver.ensureVersionTable()
41+
if err != nil {
42+
return nil, nil, xerrors.Errorf("ensure version table: %w", err)
43+
}
44+
45+
m, err := migrate.NewWithInstance("", sourceDriver, "", dbDriver)
46+
if err != nil {
47+
return nil, nil, xerrors.Errorf("new migrate instance: %w", err)
48+
}
49+
50+
return sourceDriver, m, nil
51+
}
52+
53+
type pgTxnDriver struct {
54+
ctx context.Context
55+
db *sql.DB
56+
tx *sql.Tx
57+
}
58+
59+
func (d *pgTxnDriver) Open(url string) (database.Driver, error) {
60+
panic("not implemented")
61+
}
62+
63+
func (d *pgTxnDriver) Close() error {
64+
return nil
65+
}
66+
67+
func (d *pgTxnDriver) Lock() error {
68+
var err error
69+
70+
d.tx, err = d.db.BeginTx(d.ctx, nil)
71+
if err != nil {
72+
return err
73+
}
74+
const q = `
75+
SELECT pg_advisory_xact_lock($1)
76+
`
77+
78+
_, err = d.tx.ExecContext(d.ctx, q, lockID)
79+
if err != nil {
80+
return xerrors.Errorf("exec select: %w", err)
81+
}
82+
return nil
83+
}
84+
85+
func (d *pgTxnDriver) Unlock() error {
86+
err := d.tx.Commit()
87+
d.tx = nil
88+
if err != nil {
89+
return xerrors.Errorf("commit tx on unlock: %w", err)
90+
}
91+
return nil
92+
}
93+
94+
// Run applies a migration to the database. migration is guaranteed to be not nil.
95+
func (d *pgTxnDriver) Run(migration io.Reader) error {
96+
migr, err := io.ReadAll(migration)
97+
if err != nil {
98+
return xerrors.Errorf("read migration: %w", err)
99+
}
100+
migr = bytes.ReplaceAll(migr, []byte("BEGIN;"), []byte{})
101+
migr = bytes.ReplaceAll(migr, []byte("COMMIT;"), []byte{})
102+
err = d.runStatement(migr)
103+
if err != nil {
104+
return xerrors.Errorf("run statement: %w", err)
105+
}
106+
return nil
107+
}
108+
109+
func (d *pgTxnDriver) runStatement(statement []byte) error {
110+
ctx := context.Background()
111+
query := string(statement)
112+
if strings.TrimSpace(query) == "" {
113+
return nil
114+
}
115+
if _, err := d.tx.ExecContext(ctx, query); err != nil {
116+
if pgErr, ok := err.(*pq.Error); ok {
117+
var line uint
118+
message := fmt.Sprintf("migration failed: %s", pgErr.Message)
119+
if pgErr.Detail != "" {
120+
message = fmt.Sprintf("%s, %s", message, pgErr.Detail)
121+
}
122+
return database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
123+
}
124+
return database.Error{OrigErr: err, Err: "migration failed", Query: statement}
125+
}
126+
return nil
127+
}
128+
129+
//nolint:revive
130+
func (d *pgTxnDriver) SetVersion(version int, dirty bool) error {
131+
query := `TRUNCATE ` + migrationsTableName
132+
if _, err := d.tx.Exec(query); err != nil {
133+
return &database.Error{OrigErr: err, Query: []byte(query)}
134+
}
135+
136+
if version >= 0 {
137+
query = `INSERT INTO ` + migrationsTableName + ` (version, dirty) VALUES ($1, $2)`
138+
if _, err := d.tx.Exec(query, version, dirty); err != nil {
139+
return &database.Error{OrigErr: err, Query: []byte(query)}
140+
}
141+
}
142+
143+
return nil
144+
}
145+
146+
func (d *pgTxnDriver) Version() (version int, dirty bool, err error) {
147+
var q interface {
148+
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
149+
} = d.tx
150+
if d.tx == nil {
151+
q = d.db
152+
}
153+
154+
query := `SELECT version, dirty FROM ` + migrationsTableName + ` LIMIT 1`
155+
err = q.QueryRowContext(context.Background(), query).Scan(&version, &dirty)
156+
switch {
157+
case err == sql.ErrNoRows:
158+
return database.NilVersion, false, nil
159+
160+
case err != nil:
161+
if e, ok := err.(*pq.Error); ok {
162+
if e.Code.Name() == "undefined_table" {
163+
return database.NilVersion, false, nil
164+
}
165+
}
166+
return 0, false, &database.Error{OrigErr: err, Query: []byte(query)}
167+
168+
default:
169+
return version, dirty, nil
170+
}
171+
}
172+
173+
func (d *pgTxnDriver) Drop() error {
174+
panic("not implemented")
175+
}
176+
177+
func (d *pgTxnDriver) ensureVersionTable() error {
178+
const query = `CREATE TABLE IF NOT EXISTS ` + migrationsTableName + ` (version bigint not null primary key, dirty boolean not null)`
179+
if _, err := d.db.ExecContext(context.Background(), query); err != nil {
180+
return &database.Error{OrigErr: err, Query: []byte(query)}
181+
}
182+
183+
return nil
184+
}

‎coderd/database/models.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/postgres/postgres.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ func Open() (string, func(), error) {
3636
}
3737

3838
dbName = "ci" + dbName
39-
_, err = db.Exec("CREATE DATABASE " + dbName + " WITH TEMPLATE " + os.Getenv("DB_FROM"))
39+
_, err = db.Exec("CREATE DATABASE " + dbName)
4040
if err != nil {
4141
return "", nil, xerrors.Errorf("create db with template: %w", err)
4242
}

0 commit comments

Comments
 (0)