Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Postgres: Add a check to determine if table already exists to elide CREATE query #526

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Mar 19, 2021
Merged

Postgres: Add a check to determine if table already exists to elide CREATE query #526

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
177d366
Squash commits
testtest959 Mar 16, 2021
ff67997
Format
testtest959 Mar 16, 2021
0c748e2
Minor refactoring
testtest959 Mar 16, 2021
9788837
Address PR feedback; Add mustRun
testtest959 Mar 17, 2021
8f15e2c
Fix a test assert
testtest959 Mar 18, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Squash commits
  • Loading branch information
@testtest959
testtest959 committed Mar 17, 2021
commit 177d366f8123e58a9b0f024fe5306ef1ebc4137f
19 changes: 18 additions & 1 deletion database/pgx/pgx.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,24 @@ func (p *Postgres) ensureVersionTable() (err error) {
}
}()

query := `CREATE TABLE IF NOT EXISTS ` + quoteIdentifier(p.config.MigrationsTable) + ` (version bigint not null primary key, dirty boolean not null)`
// This block checks whether the `MigrationsTable` already exists. This is useful because it allows read only postgres
// users to also check the current version of the schema. Previously, even if `MigrationsTable` existed, the
// `CREATE TABLE IF NOT EXISTS...` query would fail because the user does not have the CREATE permission.
// Taken from https://github.com/mattes/migrate/blob/master/database/postgres/postgres.go#L258
var count int
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
row := p.conn.QueryRowContext(context.Background(), query, p.config.MigrationsTable);

err = row.Scan(&count)
if err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
}

if count == 1 {
return nil
}

query = `CREATE TABLE IF NOT EXISTS ` + 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)}
}
Expand Down
154 changes: 154 additions & 0 deletions database/pgx/pgx_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"context"
"database/sql"
sqldriver "database/sql/driver"
"errors"
"fmt"
"log"

Expand Down Expand Up @@ -309,6 +310,159 @@ func TestWithSchema(t *testing.T) {
})
}

func TestFailToCreateTableWithoutPermissions(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
if err != nil {
t.Fatal(err)
}

addr := pgConnectionString(ip, port)

// Check that opening the postgres connection returns NilVersion
p := &Postgres{}

d, err := p.Open(addr)

if err != nil {
t.Fatal(err)
}

defer func() {
if err := d.Close(); err != nil {
t.Error(err)
}
}()

// create user who is not the owner. Although we're concatenating strings in an sql statement it should be fine
// since this is a test environment and we're not expecting to the pgPassword to be malicious
if err := d.Run(strings.NewReader("CREATE USER not_owner WITH ENCRYPTED PASSWORD '" + pgPassword + "'")); err != nil {
t.Fatal(err)
}

// create foobar schema
if err := d.Run(strings.NewReader("CREATE SCHEMA barfoo AUTHORIZATION postgres")); err != nil {
t.Fatal(err)
}

// revoke privileges
if err := d.Run(strings.NewReader("GRANT USAGE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM PUBLIC")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d2, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

defer func() {
if d2 == nil {
return
}
if err := d2.Close(); err != nil {
t.Fatal(err)
}
}()

var e *database.Error
if !errors.As(err, &e) || err == nil {
t.Fatal("Unexpected error, want permission denied error. Got: ", err)
}
})
}

func TestCheckBeforeCreateTable(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
if err != nil {
t.Fatal(err)
}

addr := pgConnectionString(ip, port)

// Check that opening the postgres connection returns NilVersion
p := &Postgres{}

d, err := p.Open(addr)

if err != nil {
t.Fatal(err)
}

defer func() {
if err := d.Close(); err != nil {
t.Error(err)
}
}()

// create user who is not the owner. Although we're concatenating strings in an sql statement it should be fine
// since this is a test environment and we're not expecting to the pgPassword to be malicious
if err := d.Run(strings.NewReader("CREATE USER not_owner WITH ENCRYPTED PASSWORD '" + pgPassword + "'")); err != nil {
t.Fatal(err)
}

// create foobar schema
if err := d.Run(strings.NewReader("CREATE SCHEMA barfoo AUTHORIZATION postgres")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("GRANT USAGE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("GRANT CREATE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d2, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

if err != nil {
t.Fatal(err)
}

if err := d2.Close(); err != nil {
t.Fatal(err)
}

// revoke privileges
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM PUBLIC")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d3, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

if err != nil {
t.Fatal(err)
}

version, _, err := d3.Version()

if err != nil {
t.Fatal(err)
}

if version != database.NilVersion {
t.Fatal("Unexpected version, want database.NilVersion. Got: ", version)
}

defer func() {
if err := d3.Close(); err != nil {
t.Fatal(err)
}
}()
})
}

func TestParallelSchema(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
Expand Down
19 changes: 18 additions & 1 deletion database/postgres/postgres.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,24 @@ 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)`
// This block checks whether the `MigrationsTable` already exists. This is useful because it allows read only postgres
// users to also check the current version of the schema. Previously, even if `MigrationsTable` existed, the
// `CREATE TABLE IF NOT EXISTS...` query would fail because the user does not have the CREATE permission.
// Taken from https://github.com/mattes/migrate/blob/master/database/postgres/postgres.go#L258
var count int
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
row := p.conn.QueryRowContext(context.Background(), query, p.config.MigrationsTable);

err = row.Scan(&count)
if err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
}

if count == 1 {
return nil
}

query = `CREATE TABLE IF NOT EXISTS ` + 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)}
}
Expand Down
157 changes: 155 additions & 2 deletions database/postgres/postgres_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,11 @@ import (
"context"
"database/sql"
sqldriver "database/sql/driver"
"errors"
"fmt"
"log"

"github.com/golang-migrate/migrate/v4"
"io"
"log"
"strconv"
"strings"
"sync"
Expand Down Expand Up @@ -310,6 +310,159 @@ func TestWithSchema(t *testing.T) {
})
}

func TestFailToCreateTableWithoutPermissions(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
if err != nil {
t.Fatal(err)
}

addr := pgConnectionString(ip, port)

// Check that opening the postgres connection returns NilVersion
p := &Postgres{}

d, err := p.Open(addr)

if err != nil {
t.Fatal(err)
}

defer func() {
if err := d.Close(); err != nil {
t.Error(err)
}
}()

// create user who is not the owner. Although we're concatenating strings in an sql statement it should be fine
// since this is a test environment and we're not expecting to the pgPassword to be malicious
if err := d.Run(strings.NewReader("CREATE USER not_owner WITH ENCRYPTED PASSWORD '" + pgPassword + "'")); err != nil {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like there are a bunch of setup statements for these tests. Create a mustRun(testing.T*, []string) method and use that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome! That makes this a lot cleaner

t.Fatal(err)
}

// create foobar schema
if err := d.Run(strings.NewReader("CREATE SCHEMA barfoo AUTHORIZATION postgres")); err != nil {
t.Fatal(err)
}

// revoke privileges
if err := d.Run(strings.NewReader("GRANT USAGE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM PUBLIC")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d2, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

defer func() {
if d2 == nil {
return
}
if err := d2.Close(); err != nil {
t.Fatal(err)
}
}()

var e *database.Error
if !errors.As(err, &e) || err == nil {
t.Fatal("Unexpected error, want permission denied error. Got: ", err)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the permission denied error exposed? e.g. a postgres error code

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup, but we have to parse the error string to do it. Made the change!

}
})
}

func TestCheckBeforeCreateTable(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
if err != nil {
t.Fatal(err)
}

addr := pgConnectionString(ip, port)

// Check that opening the postgres connection returns NilVersion
p := &Postgres{}

d, err := p.Open(addr)

if err != nil {
t.Fatal(err)
}

defer func() {
if err := d.Close(); err != nil {
t.Error(err)
}
}()

// create user who is not the owner. Although we're concatenating strings in an sql statement it should be fine
// since this is a test environment and we're not expecting to the pgPassword to be malicious
if err := d.Run(strings.NewReader("CREATE USER not_owner WITH ENCRYPTED PASSWORD '" + pgPassword + "'")); err != nil {
t.Fatal(err)
}

// create foobar schema
if err := d.Run(strings.NewReader("CREATE SCHEMA barfoo AUTHORIZATION postgres")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("GRANT USAGE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("GRANT CREATE ON SCHEMA barfoo TO not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d2, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

if err != nil {
t.Fatal(err)
}

if err := d2.Close(); err != nil {
t.Fatal(err)
}

// revoke privileges
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM PUBLIC")); err != nil {
t.Fatal(err)
}
if err := d.Run(strings.NewReader("REVOKE CREATE ON SCHEMA barfoo FROM not_owner")); err != nil {
t.Fatal(err)
}

// re-connect using that schema
d3, err := p.Open(fmt.Sprintf("postgres://not_owner:%s@%v:%v/postgres?sslmode=disable&search_path=barfoo",
pgPassword, ip, port))

if err != nil {
t.Fatal(err)
}

version, _, err := d3.Version()

if err != nil {
t.Fatal(err)
}

if version != database.NilVersion {
t.Fatal("Unexpected version, want database.NilVersion. Got: ", version)
}

defer func() {
if err := d3.Close(); err != nil {
t.Fatal(err)
}
}()
})
}

func TestParallelSchema(t *testing.T) {
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
ip, port, err := c.FirstPort()
Expand Down