Postgres - Add x-schema-name URL Query support (golang-migrate#95)

spacefill · Nov 28, 2020 · e800222 · e800222
1 parent a53e6fc
commit e800222
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/database/postgres/README.md b/database/postgres/README.md
@@ -5,6 +5,7 @@
 | URL Query  | WithInstance Config | Description |
 |------------|---------------------|-------------|
 | `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
+| `x-schema-name` | `SchemaName` | Schema name of the migratations table |
 | `x-statement-timeout` | `StatementTimeout` | Abort any statement that takes more than the specified number of milliseconds |
 | `dbname` | `DatabaseName` | The name of the database to connect to |
 | `search_path` | | This variable specifies the order in which schemas are searched when an object is referenced by a simple name with no schema specified. |

diff --git a/database/postgres/postgres.go b/database/postgres/postgres.go
@@ -123,6 +123,7 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
 	}
 
 	migrationsTable := purl.Query().Get("x-migrations-table")
+	schemaName := purl.Query().Get("x-schema-name")
 	statementTimeoutString := purl.Query().Get("x-statement-timeout")
 	statementTimeout := 0
 	if statementTimeoutString != "" {
@@ -135,6 +136,7 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
 	px, err := WithInstance(db, &Config{
 		DatabaseName:     purl.Path,
 		MigrationsTable:  migrationsTable,
+		SchemaName:       schemaName,
 		StatementTimeout: time.Duration(statementTimeout) * time.Millisecond,
 	})
 
@@ -272,7 +274,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
 		return &database.Error{OrigErr: err, Err: "transaction start failed"}
 	}
 
-	query := `TRUNCATE ` + pq.QuoteIdentifier(p.config.MigrationsTable)
+	query := `TRUNCATE ` + pq.QuoteIdentifier(p.config.SchemaName) + `.` + pq.QuoteIdentifier(p.config.MigrationsTable)
 	if _, err := tx.Exec(query); err != nil {
 		if errRollback := tx.Rollback(); errRollback != nil {
 			err = multierror.Append(err, errRollback)
@@ -284,7 +286,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
 	// empty schema version for failed down migration on the first migration
 	// See: https://github.com/golang-migrate/migrate/issues/330
 	if version >= 0 || (version == database.NilVersion && dirty) {
-		query = `INSERT INTO ` + pq.QuoteIdentifier(p.config.MigrationsTable) +
+		query = `INSERT INTO ` + pq.QuoteIdentifier(p.config.SchemaName) + `.` + pq.QuoteIdentifier(p.config.MigrationsTable) +
 			` (version, dirty) VALUES ($1, $2)`
 		if _, err := tx.Exec(query, version, dirty); err != nil {
 			if errRollback := tx.Rollback(); errRollback != nil {
@@ -302,7 +304,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
 }
 
 func (p *Postgres) Version() (version int, dirty bool, err error) {
-	query := `SELECT version, dirty FROM ` + pq.QuoteIdentifier(p.config.MigrationsTable) + ` LIMIT 1`
+	query := `SELECT version, dirty FROM ` + pq.QuoteIdentifier(p.config.SchemaName) + `.` + pq.QuoteIdentifier(p.config.MigrationsTable) + ` LIMIT 1`
 	err = p.conn.QueryRowContext(context.Background(), query).Scan(&version, &dirty)
 	switch {
 	case err == sql.ErrNoRows:
@@ -380,7 +382,7 @@ func (p *Postgres) ensureVersionTable() (err error) {
 		}
 	}()
 
-	query := `CREATE TABLE IF NOT EXISTS ` + pq.QuoteIdentifier(p.config.MigrationsTable) + ` (version bigint not null primary key, dirty boolean not null)`
+	query := `CREATE TABLE IF NOT EXISTS ` + pq.QuoteIdentifier(p.config.SchemaName) + `.` + pq.QuoteIdentifier(p.config.MigrationsTable) + ` (version bigint not null primary key, dirty boolean not null)`
 	if _, err = p.conn.ExecContext(context.Background(), query); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}