Skip to content

Commit cf2fbf4

Browse files
stephane-kleinstephane-klein
committed
Postgres and pgx - Add x-migrations-table-quoted url query option to postgres and pgx drivers (#95)
By default, gomigrate quote migrations table name, if `x-migrations-table-quoted` is enabled, then you must to quote migrations table name manually, for instance `"gomigrate"."schema_migrations"`
1 parent 9e969a0 commit cf2fbf4

File tree

6 files changed

+309
-20
lines changed

6 files changed

+309
-20
lines changed

‎database/pgx/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
| URL Query | WithInstance Config | Description |
66
|------------|---------------------|-------------|
77
| `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
8+
| `x-migrations-table-quoted` | `MigrationsTableQuoted` | By default, gomigrate quote migrations table name, if `x-migrations-table-quoted` is enabled, then you must quote migrations table name manually, for instance `"gomigrate"."schema_migrations"` |
89
| `x-statement-timeout` | `StatementTimeout` | Abort any statement that takes more than the specified number of milliseconds |
910
| `x-multi-statement` | `MultiStatementEnabled` | Enable multi-statement execution (default: false) |
1011
| `x-multi-statement-max-size` | `MultiStatementMaxSize` | Maximum size of single statement in bytes (default: 10MB) |

‎database/pgx/pgx.go‎

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"io"
1010
"io/ioutil"
1111
nurl "net/url"
12+
"regexp"
1213
"strconv"
1314
"strings"
1415
"time"
@@ -46,6 +47,7 @@ type Config struct {
4647
DatabaseName string
4748
SchemaName string
4849
StatementTimeout time.Duration
50+
MigrationsTableQuoted bool
4951
MultiStatementEnabled bool
5052
MultiStatementMaxSize int
5153
}
@@ -137,6 +139,17 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
137139
}
138140

139141
migrationsTable := purl.Query().Get("x-migrations-table")
142+
migrationsTableQuoted := false
143+
if s := purl.Query().Get("x-migrations-table-quoted"); len(s) > 0 {
144+
migrationsTableQuoted, err = strconv.ParseBool(s)
145+
if err != nil {
146+
return nil, fmt.Errorf("Unable to parse option x-migrations-table-quoted: %w", err)
147+
}
148+
}
149+
if (len(migrationsTable) > 0) && (migrationsTableQuoted) && ((migrationsTable[0] != '"') || (migrationsTable[len(migrationsTable)-1] != '"')) {
150+
return nil, fmt.Errorf("x-migrations-table must be quoted (for instance '\"gomigrate\".\"schema_migrations\"') when x-migrations-table-quoted is enabled, current value is: %s", migrationsTable)
151+
}
152+
140153
statementTimeoutString := purl.Query().Get("x-statement-timeout")
141154
statementTimeout := 0
142155
if statementTimeoutString != "" {
@@ -168,6 +181,7 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
168181
px, err := WithInstance(db, &Config{
169182
DatabaseName: purl.Path,
170183
MigrationsTable: migrationsTable,
184+
MigrationsTableQuoted: migrationsTableQuoted,
171185
StatementTimeout: time.Duration(statementTimeout) * time.Millisecond,
172186
MultiStatementEnabled: multiStatementEnabled,
173187
MultiStatementMaxSize: multiStatementMaxSize,
@@ -321,7 +335,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
321335
return &database.Error{OrigErr: err, Err: "transaction start failed"}
322336
}
323337

324-
query := `TRUNCATE ` + quoteIdentifier(p.config.MigrationsTable)
338+
query := `TRUNCATE ` + p.quoteIdentifier(p.config.MigrationsTable)
325339
if _, err := tx.Exec(query); err != nil {
326340
if errRollback := tx.Rollback(); errRollback != nil {
327341
err = multierror.Append(err, errRollback)
@@ -333,7 +347,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
333347
// empty schema version for failed down migration on the first migration
334348
// See: https://github.com/golang-migrate/migrate/issues/330
335349
if version >= 0 || (version == database.NilVersion && dirty) {
336-
query = `INSERT INTO ` + quoteIdentifier(p.config.MigrationsTable) +
350+
query = `INSERT INTO ` + p.quoteIdentifier(p.config.MigrationsTable) +
337351
` (version, dirty) VALUES ($1, $2)`
338352
if _, err := tx.Exec(query, version, dirty); err != nil {
339353
if errRollback := tx.Rollback(); errRollback != nil {
@@ -351,7 +365,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
351365
}
352366

353367
func (p *Postgres) Version() (version int, dirty bool, err error) {
354-
query := `SELECT version, dirty FROM ` + quoteIdentifier(p.config.MigrationsTable) + ` LIMIT 1`
368+
query := `SELECT version, dirty FROM ` + p.quoteIdentifier(p.config.MigrationsTable) + ` LIMIT 1`
355369
err = p.conn.QueryRowContext(context.Background(), query).Scan(&version, &dirty)
356370
switch {
357371
case err == sql.ErrNoRows:
@@ -401,7 +415,7 @@ func (p *Postgres) Drop() (err error) {
401415
if len(tableNames) > 0 {
402416
// delete one by one ...
403417
for _, t := range tableNames {
404-
query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
418+
query = `DROP TABLE IF EXISTS ` + p.quoteIdentifier(t) + ` CASCADE`
405419
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
406420
return &database.Error{OrigErr: err, Query: []byte(query)}
407421
}
@@ -433,10 +447,27 @@ func (p *Postgres) ensureVersionTable() (err error) {
433447
// users to also check the current version of the schema. Previously, even if `MigrationsTable` existed, the
434448
// `CREATE TABLE IF NOT EXISTS...` query would fail because the user does not have the CREATE permission.
435449
// Taken from https://github.com/mattes/migrate/blob/master/database/postgres/postgres.go#L258
436-
var count int
437-
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
438-
row := p.conn.QueryRowContext(context.Background(), query, p.config.MigrationsTable)
450+
var row *sql.Row
451+
tableName := []byte(p.config.MigrationsTable)
452+
schemaName := []byte("")
453+
if p.config.MigrationsTableQuoted {
454+
re := regexp.MustCompile(`"(.*?)"`)
455+
result := re.FindAllSubmatch([]byte(p.config.MigrationsTable), -1)
456+
tableName = result[len(result)-1][1]
457+
if len(result) > 1 {
458+
schemaName = result[0][1]
459+
}
460+
}
461+
var query string
462+
if len(schemaName) > 0 {
463+
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = $2 LIMIT 1`
464+
row = p.conn.QueryRowContext(context.Background(), query, tableName, schemaName)
465+
} else {
466+
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
467+
row = p.conn.QueryRowContext(context.Background(), query, tableName)
468+
}
439469

470+
var count int
440471
err = row.Scan(&count)
441472
if err != nil {
442473
return &database.Error{OrigErr: err, Query: []byte(query)}
@@ -446,7 +477,7 @@ func (p *Postgres) ensureVersionTable() (err error) {
446477
return nil
447478
}
448479

449-
query = `CREATE TABLE IF NOT EXISTS ` + quoteIdentifier(p.config.MigrationsTable) + ` (version bigint not null primary key, dirty boolean not null)`
480+
query = `CREATE TABLE IF NOT EXISTS ` + p.quoteIdentifier(p.config.MigrationsTable) + ` (version bigint not null primary key, dirty boolean not null)`
450481
if _, err = p.conn.ExecContext(context.Background(), query); err != nil {
451482
return &database.Error{OrigErr: err, Query: []byte(query)}
452483
}
@@ -455,7 +486,10 @@ func (p *Postgres) ensureVersionTable() (err error) {
455486
}
456487

457488
// Copied from lib/pq implementation: https://github.com/lib/pq/blob/v1.9.0/conn.go#L1611
458-
func quoteIdentifier(name string) string {
489+
func (p *Postgres) quoteIdentifier(name string) string {
490+
if p.config.MigrationsTableQuoted {
491+
return name
492+
}
459493
end := strings.IndexRune(name, 0)
460494
if end > -1 {
461495
name = name[:end]

‎database/pgx/pgx_test.go‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,65 @@ func TestWithSchema(t *testing.T) {
318318
})
319319
}
320320

321+
func TestMigrationTableOption(t *testing.T) {
322+
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
323+
ip, port, err := c.FirstPort()
324+
if err != nil {
325+
t.Fatal(err)
326+
}
327+
328+
addr := pgConnectionString(ip, port)
329+
p := &Postgres{}
330+
d, _ := p.Open(addr)
331+
defer func() {
332+
if err := d.Close(); err != nil {
333+
t.Fatal(err)
334+
}
335+
}()
336+
337+
// create gomigrate schema
338+
if err := d.Run(strings.NewReader("CREATE SCHEMA gomigrate AUTHORIZATION postgres")); err != nil {
339+
t.Fatal(err)
340+
}
341+
342+
// bad unquoted x-migrations-table parameter
343+
d, err = p.Open(fmt.Sprintf("postgres://postgres:%s@%v:%v/postgres?sslmode=disable&x-migrations-table=gomigrate.schema_migrations&x-migrations-table-quoted=1",
344+
pgPassword, ip, port))
345+
if err == nil {
346+
t.Fatal("expected x-migrations-table must be quoted...")
347+
}
348+
349+
// good quoted x-migrations-table parameter
350+
d, err = p.Open(fmt.Sprintf("postgres://postgres:%s@%v:%v/postgres?sslmode=disable&x-migrations-table=\"gomigrate\".\"schema_migrations\"&x-migrations-table-quoted=1",
351+
pgPassword, ip, port))
352+
if err != nil {
353+
t.Fatal(err)
354+
}
355+
356+
// make sure gomigrate.schema_migrations table exists
357+
var exists bool
358+
if err := d.(*Postgres).conn.QueryRowContext(context.Background(), "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'schema_migrations' AND table_schema = 'gomigrate')").Scan(&exists); err != nil {
359+
t.Fatal(err)
360+
}
361+
if !exists {
362+
t.Fatalf("expected table gomigrate.schema_migrations to exist")
363+
}
364+
365+
d, err = p.Open(fmt.Sprintf("postgres://postgres:%s@%v:%v/postgres?sslmode=disable&x-migrations-table=gomigrate.schema_migrations",
366+
pgPassword, ip, port))
367+
if err != nil {
368+
t.Fatal(err)
369+
}
370+
if err := d.(*Postgres).conn.QueryRowContext(context.Background(), "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'gomigrate.schema_migrations' AND table_schema = (SELECT current_schema()))").Scan(&exists); err != nil {
371+
t.Fatal(err)
372+
}
373+
if !exists {
374+
t.Fatalf("expected table 'gomigrate.schema_migrations' to exist")
375+
}
376+
377+
})
378+
}
379+
321380
func TestFailToCreateTableWithoutPermissions(t *testing.T) {
322381
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
323382
ip, port, err := c.FirstPort()
@@ -373,6 +432,18 @@ func TestFailToCreateTableWithoutPermissions(t *testing.T) {
373432
if !strings.Contains(e.OrigErr.Error(), "permission denied for schema barfoo") {
374433
t.Fatal(e)
375434
}
435+
436+
// re-connect using that x-migrations-table and x-migrations-table-quoted
437+
d2, err = p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&x-migrations-table=\"barfoo\".\"schema_migrations\"&x-migrations-table-quoted=1",
438+
pgPassword, ip, port))
439+
440+
if !errors.As(err, &e) || err == nil {
441+
t.Fatal("Unexpected error, want permission denied error. Got: ", err)
442+
}
443+
444+
if !strings.Contains(e.OrigErr.Error(), "permission denied for schema barfoo") {
445+
t.Fatal(e)
446+
}
376447
})
377448
}
378449

@@ -679,5 +750,41 @@ func Test_computeLineFromPos(t *testing.T) {
679750
run(true, true)
680751
})
681752
}
753+
}
682754

755+
func Test_quoteIdentifier(t *testing.T) {
756+
testcases := []struct {
757+
migrationsTableQuoted bool
758+
migrationsTable string
759+
want string
760+
}{
761+
{
762+
false,
763+
"schema_name.table_name",
764+
"\"schema_name.table_name\"",
765+
},
766+
{
767+
false,
768+
"table_name",
769+
"\"table_name\"",
770+
},
771+
{
772+
true,
773+
"\"schema_name\".\"table.name\"",
774+
"\"schema_name\".\"table.name\"",
775+
},
776+
}
777+
p := &Postgres{
778+
config: &Config{
779+
MigrationsTableQuoted: false,
780+
},
781+
}
782+
783+
for _, tc := range testcases {
784+
p.config.MigrationsTableQuoted = tc.migrationsTableQuoted
785+
got := p.quoteIdentifier(tc.migrationsTable)
786+
if tc.want != got {
787+
t.Fatalf("expected %s but got %s", tc.want, got)
788+
}
789+
}
683790
}

‎database/postgres/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
| URL Query | WithInstance Config | Description |
66
|------------|---------------------|-------------|
77
| `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
8+
| `x-migrations-table-quoted` | `MigrationsTableQuoted` | By default, gomigrate quote migrations table name, if `x-migrations-table-quoted` is enabled, then you must quote migrations table name manually, for instance `"gomigrate"."schema_migrations"` |
89
| `x-statement-timeout` | `StatementTimeout` | Abort any statement that takes more than the specified number of milliseconds |
910
| `x-multi-statement` | `MultiStatementEnabled` | Enable multi-statement execution (default: false) |
1011
| `x-multi-statement-max-size` | `MultiStatementMaxSize` | Maximum size of single statement in bytes (default: 10MB) |

0 commit comments

Comments
 (0)