From 1e2fe258b90fe2d1b9634da473688e6818509cf9 Mon Sep 17 00:00:00 2001 From: Peter Downs Date: Fri, 29 Mar 2024 11:42:46 -0400 Subject: [PATCH] feat: support custom roles for superuser-only extensions like postgis --- .github/workflows/golang.yaml | 2 +- README.md | 94 +++++++++++++++++---- docker-compose.yml | 6 +- testdb.go | 149 ++++++++++++++++++++++++++-------- testdb_test.go | 88 ++++++++++++++++++++ 5 files changed, 288 insertions(+), 51 deletions(-) diff --git a/.github/workflows/golang.yaml b/.github/workflows/golang.yaml index e600088..bfa298a 100644 --- a/.github/workflows/golang.yaml +++ b/.github/workflows/golang.yaml @@ -8,7 +8,7 @@ jobs: test: services: postgres: - image: postgres:15 + image: postgis/postgis:15-master env: POSTGRES_PASSWORD: password # TODO: unable to turn off fsync easily, see diff --git a/README.md b/README.md index e39cb9a..f61867f 100644 --- a/README.md +++ b/README.md @@ -12,23 +12,26 @@ everything. # Documentation -- [This page, https://github.com/peterldowns/pgtestdb](https://github.com/peterldowns/pgtestdb) +- [The github README, https://github.com/peterldowns/pgtestdb](https://github.com/peterldowns/pgtestdb) - [The go.dev docs, pkg.go.dev/github.com/peterldowns/pgtestdb](https://pkg.go.dev/github.com/peterldowns/pgtestdb) -This page is the primary source for documentation. The code itself is supposed -to be well-organized, and each function has a meaningful docstring, so you -should be able to explore it quite easily using an LSP plugin, reading the +The github README is the primary source for documentation. The code itself is +supposed to be well-organized, and each function has a meaningful docstring, so +you should be able to explore it quite easily using an LSP plugin, reading the code, or clicking through the go.dev docs. ## How does it work? -Each time a test asks for a fresh database by calling `pgtestdb.New`, pgtestdb will +Each time one of your tests asks for a fresh database by calling `pgtestdb.New`, pgtestdb will check to see if a template database already exists. If not, it creates a new -database, runs your migrations on it, and then marks it as a template. Once the -template exists, it is _very_ fast to create a new database from that template. +database, runs your migrations on it, and then marks it as a template. Once the +template exists, it then creates a test-specific database from that template. -pgtestdb only runs migrations one time when your migrations change. The marginal -cost of a new test that uses the database is just the time to create a clone -from the template, which is now basically free. +Creating a new database from a template is _very_ fast, on the order of 10s of +milliseconds. And because pgtestdb uses advisory locks and hashes your +migrations to determine which template database to use, your migrations only +end up being run one time, regardless of how many tests or separate packages +you have. This is true even across test runs --- pgtestdb will only run your +migrations again if you change them in some way. When a test succeeds, the database it used is automatically deleted. When a test fails, the database it used is left alive, and the test logs will @@ -334,13 +337,18 @@ config (see below.) ```go // Config contains the details needed to connect to a postgres server/database. type Config struct { - DriverName string // "pgx" (pgx) or "postgres" (lib/pq) - Host string // "localhost" - Port string // "5433" - User string // "postgres" - Password string // "password" - Database string // "postgres" - Options string // "sslmode=disable&anotherSetting=value" + DriverName string // the name of a driver to use when calling sql.Open() to connect to a database, "pgx" (pgx) or "postgres" (lib/pq) + Host string // the host of the database, "localhost" + Port string // the port of the database, "5433" + User string // the user to connect as, "postgres" + Password string // the password to connect with, "password" + Database string // the database to connect to, "postgres" + Options string // URL-formatted additional options to pass in the connection string, "sslmode=disable&something=value" + // TestRole is the role used to create and connect to the template database + // and each test database. If not provided, defaults to [DefaultRole]. The + // capabilities of this role should match the capabilities of the role that + // your application uses to connect to its database and run migrations. + TestRole *Role } // URL returns a postgres connection string in the format @@ -357,6 +365,58 @@ new databases and roles. Most likely you want to connect as the default `postgres` user, since you'll be connecting to a dedicated testing-only Postgres server as described earlier. +### `pgtestdb.Role` +A dedicated Postgres role (user) is used to create the template database and each test database. pgtestdb will create this role for you with sane defaults, but you can control the username, password, and capabilities of this role if desired. + +```go +const ( + // DefaultRoleUsername is the default name for the role that is created and + // used to create and connect to each test database. + DefaultRoleUsername = "pgtdbuser" + // DefaultRolePassword is the default password for the role that is created and + // used to create and connect to each test database. + DefaultRolePassword = "pgtdbpass" + // DefaultRoleCapabilities is the default set of capabilities for the role + // that is created and used to create and conect to each test database. + // This is locked down by default, and will not allow the creation of + // extensions. + DefaultRoleCapabilities = "NOSUPERUSER NOCREATEDB NOCREATEROLE" +) + +// DefaultRole returns the default Role used to create and connect to the +// template database and each test database. It is a function, not a struct, to +// prevent accidental overriding. +func DefaultRole() Role { + return Role{ + Username: DefaultRoleUsername, + Password: DefaultRolePassword, + Capabilities: DefaultRoleCapabilities, + } +} + +// Role contains the details of a postgres role (user) that will be used +// when creating and connecting to the template and test databases. +type Role struct { + // The username for the role, defaults to [DefaultRoleUsername]. + Username string + // The password for the role, defaults to [DefaultRolePassword]. + Password string + // The capabilities that will be granted to the role, defaults to + // [DefaultRoleCapabilities]. + Capabilities string +} +``` + +Because this role is used to connect to each template and each test database +and run the migrations, its capabilities should match those of your production +application. For instance, if in production your application connects as a +superuser, you will want to pass a custom `Role` whthat includes the +`SUPERUSER` capability so that your migrations will run the same in both +envproduction and tests. + +This is a common case for many applications that install or activate extensions +like [Postgis](https://postgis.net/), which require activation via a superuser. + ### `pgtestdb.Migrator` The `Migrator` interface contains all of the logic needed to prepare a template diff --git a/docker-compose.yml b/docker-compose.yml index 3b7411e..32931d4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,7 +1,11 @@ version: "3.6" services: testdb: - image: postgres:15 + # We're using postgis so that we can test the creation of the postgis + # extension, which requires superuser extensions. + # + # To use the equivalent in plain postgres, use `postgres:15` + image: postgis/postgis:15-master environment: POSTGRES_PASSWORD: password restart: unless-stopped diff --git a/testdb.go b/testdb.go index f1819f8..c08ce63 100644 --- a/testdb.go +++ b/testdb.go @@ -7,28 +7,67 @@ import ( "database/sql" "encoding/hex" "fmt" - "testing" "github.com/peterldowns/pgtestdb/internal/once" "github.com/peterldowns/pgtestdb/internal/sessionlock" + "github.com/peterldowns/pgtestdb/migrators/common" ) const ( - // TestUser is the username for connecting to each test database. - TestUser = "pgtdbuser" - // TestPassword is the password for connecting to each test database. - TestPassword = "pgtdbpass" + // DefaultRoleUsername is the default name for the role that is created and + // used to create and connect to each test database. + DefaultRoleUsername = "pgtdbuser" + // DefaultRolePassword is the default password for the role that is created and + // used to create and connect to each test database. + DefaultRolePassword = "pgtdbpass" + // DefaultRoleCapabilities is the default set of capabilities for the role + // that is created and used to create and conect to each test database. + // This is locked down by default, and will not allow the creation of + // extensions. + DefaultRoleCapabilities = "NOSUPERUSER NOCREATEDB NOCREATEROLE" + // Deprecated: prefer [DefaultRoleUsername]. + TestUser = DefaultRoleUsername + // Deprecated: prefer [DefaultRolePassword]. + TestPassword = DefaultRolePassword ) +// DefaultRole returns the default Role used to create and connect to the +// template database and each test database. It is a function, not a struct, to +// prevent accidental overriding. +func DefaultRole() Role { + return Role{ + Username: DefaultRoleUsername, + Password: DefaultRolePassword, + Capabilities: DefaultRoleCapabilities, + } +} + // Config contains the details needed to connect to a postgres server/database. type Config struct { - DriverName string // the name of a driver to use when calling sql.Open() to connect to a database + DriverName string // the name of a driver to use when calling sql.Open() to connect to a database, "pgx" (pgx) or "postgres" (lib/pq) Host string // the host of the database, "localhost" Port string // the port of the database, "5433" User string // the user to connect as, "postgres" Password string // the password to connect with, "password" Database string // the database to connect to, "postgres" Options string // URL-formatted additional options to pass in the connection string, "sslmode=disable&something=value" + // TestRole is the role used to create and connect to the template database + // and each test database. If not provided, defaults to [DefaultRole]. The + // capabilities of this role should match the capabilities of the role that + // your application uses to connect to its database and run migrations. + TestRole *Role +} + +// Role contains the details of a postgres role (user) that will be used +// when creating and connecting to the template and test databases. +type Role struct { + // The username for the role, defaults to [DefaultRoleUsername]. + Username string + // The password for the role, defaults to [DefaultRolePassword]. + Password string + // The capabilities that will be granted to the role, defaults to + // [DefaultRoleCapabilities]. + Capabilities string } // URL returns a postgres connection string in the format @@ -87,7 +126,7 @@ type Migrator interface { // database instance. This database is prepared and migrated by the given // migrator, by get-or-creating a template database and then cloning it. This is // a concurrency-safe primitive. If there is an error creating the database, the -// test will be immediately failed with `t.Fatal()`. +// test will be immediately failed with `t.Fatalf()`. // // If this method succeeds, it will `t.Log()` the connection string to the // created database, so that if your test fails, you can connect to the database @@ -96,20 +135,32 @@ type Migrator interface { // If this method succeeds and your test succeeds, the database will be removed // as part of the test cleanup process. // -// `testing.TB` is the common testing interface -// implemented by `*testing.T`, `*testing.B`, and `*testing.F`, so you can use -// pgtestdb to get a database for tests, benchmarks, and fuzzes. -func New(t testing.TB, conf Config, migrator Migrator) *sql.DB { +// `TB` is a subset of the `testing.TB` testing interface implemented by +// `*testing.T`, `*testing.B`, and `*testing.F`, so you can use pgtestdb to get +// a database for tests, benchmarks, and fuzzes. +func New(t TB, conf Config, migrator Migrator) *sql.DB { t.Helper() _, db := create(t, conf, migrator) return db } +// TB is a subset of the `testing.TB` testing interface implemented by +// `*testing.T`, `*testing.B`, and `*testing.F`, so you can use pgtestdb to get +// a database for tests, benchmarks, and fuzzes. It contains only the methods +// actually needed by pgtestdb, defined so that we can more easily mock it. +type TB interface { + Cleanup(func()) + Failed() bool + Fatalf(format string, args ...any) + Helper() + Logf(format string, args ...any) +} + // Custom is like [New] but after creating the new database instance, it closes // any connections and returns the configuration details of that database so // that you can connect to it explicitly, potentially via a different SQL // interface. -func Custom(t testing.TB, conf Config, migrator Migrator) *Config { +func Custom(t TB, conf Config, migrator Migrator) *Config { t.Helper() config, db := create(t, conf, migrator) // Close `*sql.DB` connection that was opened during the creation process so @@ -122,7 +173,14 @@ func Custom(t testing.TB, conf Config, migrator Migrator) *Config { return config } -func create(t testing.TB, conf Config, migrator Migrator) (*Config, *sql.DB) { +// Helper +// Fatalf +// Fatal +// Logf +// Cleanup +// Failed + +func create(t TB, conf Config, migrator Migrator) (*Config, *sql.DB) { t.Helper() ctx := context.Background() baseDB, err := conf.Connect() @@ -131,13 +189,21 @@ func create(t testing.TB, conf Config, migrator Migrator) (*Config, *sql.DB) { return nil, nil // unreachable } - if err := ensureUser(ctx, baseDB); err != nil { + // From this point onward, all functions assume that `conf.TestRole` is not nil. + // We default to the + if conf.TestRole == nil { + role := DefaultRole() + conf.TestRole = &role + } + if err := ensureUser(ctx, baseDB, conf); err != nil { t.Fatalf("could not create pgtestdb user: %s", err) + return nil, nil // unreachable } template, err := getOrCreateTemplate(ctx, baseDB, conf, migrator) if err != nil { - t.Fatal(err) + t.Fatalf("%s", err) + return nil, nil // unreachable } instance, err := createInstance(ctx, baseDB, *template) @@ -167,6 +233,7 @@ func create(t testing.TB, conf Config, migrator Migrator) (*Config, *sql.DB) { query := fmt.Sprintf(`DROP DATABASE IF EXISTS "%s"`, instance.Database) if _, err := baseDB.ExecContext(ctx, query); err != nil { t.Fatalf("could not drop test database '%s': %s", instance.Database, err) + return // unreachable } }) @@ -179,43 +246,50 @@ func create(t testing.TB, conf Config, migrator Migrator) (*Config, *sql.DB) { // Assumption: verification is >>> faster than performing the migrations, // and is therefore safe to run at the beginning of each test. if err := migrator.Verify(ctx, db, *instance); err != nil { - t.Fatal(fmt.Errorf("test database failed verification %s: %w", instance.Database, err)) + t.Fatalf("test database failed verification %s: %w", instance.Database, err) + return nil, nil // unreachable } return instance, db } -// user is used to guarantee that the testdb user/role is only get-or-created at -// most once per program. -var user once.Var[any] = once.NewVar[any]() //nolint:gochecknoglobals +// user is used to guarantee that each testdb user/role is only get-or-created +// at most once per program. Different calls to pgtestdb can specify different +// roles, but each will be get-or-created at most one time per program, and will +// be created only once no matter how many different programs or test suites run +// at once, thanks to the use of session locks. +var users once.Map[string, any] = once.NewMap[string, any]() //nolint:gochecknoglobals func ensureUser( ctx context.Context, baseDB *sql.DB, + conf Config, ) error { - _, err := user.Set(func() (*any, error) { - return nil, sessionlock.With(ctx, baseDB, "testdb-user", func(conn *sql.Conn) error { + username := conf.TestRole.Username + _, err := users.Set(username, func() (*any, error) { + return nil, sessionlock.With(ctx, baseDB, username, func(conn *sql.Conn) error { // Get-or-create a role/user dedicated to connecting to these test databases. var roleExists bool query := "SELECT EXISTS (SELECT from pg_catalog.pg_roles WHERE rolname = $1)" - if err := conn.QueryRowContext(ctx, query, TestUser).Scan(&roleExists); err != nil { - return fmt.Errorf("failed to detect if role %s exists: %w", TestUser, err) + if err := conn.QueryRowContext(ctx, query, username).Scan(&roleExists); err != nil { + return fmt.Errorf("failed to detect if role %s exists: %w", username, err) } if roleExists { return nil } if !roleExists { - query = fmt.Sprintf(`CREATE ROLE "%s"`, TestUser) + query = fmt.Sprintf(`CREATE ROLE "%s"`, username) if _, err := conn.ExecContext(ctx, query); err != nil { - return fmt.Errorf("failed to create role %s: %w", TestUser, err) + return fmt.Errorf("failed to create role %s: %w", username, err) } query = fmt.Sprintf( - `ALTER ROLE "%s" WITH LOGIN PASSWORD '%s' NOSUPERUSER NOCREATEDB NOCREATEROLE`, - TestUser, - TestPassword, + `ALTER ROLE "%s" WITH LOGIN PASSWORD '%s' %s`, + username, + conf.TestRole.Password, + conf.TestRole.Capabilities, ) if _, err := conn.ExecContext(ctx, query); err != nil { - return fmt.Errorf("failed to alter role and set password for %s: %w", TestUser, err) + return fmt.Errorf("failed to set password and capabilities for '%s': %w", username, err) } } return nil @@ -252,10 +326,20 @@ func getOrCreateTemplate( dbconf Config, migrator Migrator, ) (*templateState, error) { - hash, err := migrator.Hash() + mhash, err := migrator.Hash() if err != nil { return nil, fmt.Errorf("failed to calculate template hash: %w", err) } + // The migrator Hash() implementation is included, along with the role + // details, so that if the user runs tests in parallel with different role + // information, they each get their own database. + hash := common.NewRecursiveHash( + common.Field("Username", dbconf.TestRole.Username), + common.Field("Password", dbconf.TestRole.Password), + common.Field("Capabilities", dbconf.TestRole.Capabilities), + common.Field("MigratorHash", mhash), + ).String() + return states.Set(hash, func() (*templateState, error) { // This function runs once per program, but only synchronizes access // within a single program. When running larger test suites, each @@ -264,8 +348,9 @@ func getOrCreateTemplate( state := templateState{} state.hash = hash state.conf = dbconf - state.conf.User = TestUser - state.conf.Password = TestPassword + state.conf.TestRole = dbconf.TestRole + state.conf.User = dbconf.TestRole.Username + state.conf.Password = dbconf.TestRole.Password state.conf.Database = fmt.Sprintf("testdb_tpl_%s", hash) // sessionlock synchronizes the creation of the template with a // session-scoped advisory lock. diff --git a/testdb_test.go b/testdb_test.go index 848365a..2fe3e90 100644 --- a/testdb_test.go +++ b/testdb_test.go @@ -238,6 +238,61 @@ func TestMigrationWithConcurrentCreate(t *testing.T) { } } +func TestDefaultRolePreventsPostgis(t *testing.T) { + t.Parallel() + config := pgtestdb.Config{ + DriverName: "pgx", + User: "postgres", + Password: "password", + Host: "localhost", + Port: "5433", + Options: "sslmode=disable", + } + migrator := &sqlMigrator{ + migrations: []string{ + // This requires SUPERUSER permissions, but by default they're + // not enabled, so it should fail. + "CREATE EXTENSION postgis;", + }, + } + tt := &MockT{} + _ = pgtestdb.New(tt, config, migrator) + tt.DoCleanup() + assert.True(t, tt.Failed()) +} + +func TestCustomRoleAllowsPostgis(t *testing.T) { + t.Parallel() + config := pgtestdb.Config{ + DriverName: "pgx", + User: "postgres", + Password: "password", + Host: "localhost", + Port: "5433", + Options: "sslmode=disable", + TestRole: &pgtestdb.Role{ + // Must use a distinct name or it will collide with other tests that + // use the default username, but have non-SUPERUSER capabilities. + // TODO: figure out some way to detect a difference in capabilities + // and fail + warn the user about the collision. + // Or, TODO: figure out a way to include a hash of the capabilities + // on to the basename if there are custom capabilities? But then + // it's a pain and confusing. Blargh. + Username: "pgtestdb-superuser", + Password: pgtestdb.DefaultRolePassword, + Capabilities: "SUPERUSER", + }, + } + migrator := &sqlMigrator{ + migrations: []string{ + // This will work since the migrations will be run with a role that + // has SUPERUSER permissions. + "CREATE EXTENSION postgis;", + }, + } + _ = pgtestdb.New(t, config, migrator) +} + // pgtestdb.New should be able to connect with either lib/pq or pgx/stdlib. func TestWithLibPqAndPgxStdlibDrivers(t *testing.T) { t.Parallel() @@ -349,3 +404,36 @@ func (s *sqlMigrator) Verify(ctx context.Context, db *sql.DB, _ pgtestdb.Config) } return nil } + +// MockT implements the `TB“ interface so that we can check to see if a test +// "would have failed". +type MockT struct { + failed bool + cleanups []func() +} + +func (t *MockT) Fatalf(string, ...any) { + t.failed = true +} + +func (*MockT) Logf(string, ...any) { + // no-op +} + +func (*MockT) Helper() { + // no-op +} + +func (t *MockT) Cleanup(f func()) { + t.cleanups = append(t.cleanups, f) +} + +func (t *MockT) DoCleanup() { + for _, f := range t.cleanups { + f() + } +} + +func (t *MockT) Failed() bool { + return t.failed +}