Skip to content
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.

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

feat(postgres): use faster sql.DB instead of docker exec psql for snapshot/restore [rebased for main] #2600

Merged
Merged
Changes from 1 commit
Commits
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
feat(postgres): use faster sql.DB instead of docker exec psql for c…
…reating & restoring snapshots
cfstras committed Jun 21, 2024
commit 51f83fffd9dad86c8e49cf13f047dacc0feaba66
111 changes: 73 additions & 38 deletions modules/postgres/postgres.go
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@ package postgres

import (
"context"
"database/sql"
"fmt"
"io"
"net"
@@ -187,54 +188,45 @@ func WithSnapshotName(name string) SnapshotOption {
// customize the snapshot name with the options.
// If a snapshot already exists under the given/default name, it will be overwritten with the new snapshot.
func (c *PostgresContainer) Snapshot(ctx context.Context, opts ...SnapshotOption) error {
config := &snapshotConfig{}
for _, opt := range opts {
config = opt(config)
}

snapshotName := defaultSnapshotName
if config.snapshotName != "" {
snapshotName = config.snapshotName
}

if c.dbName == "postgres" {
return fmt.Errorf("cannot snapshot the postgres system database as it cannot be dropped to be restored")
snapshotName, err := c.checkSnapshotConfig(opts)
if err != nil {
return err
}

// execute the commands to create the snapshot, in order
cmds := []string{
if err := c.execCommandsSQL(ctx,
// Drop the snapshot database if it already exists
fmt.Sprintf(`DROP DATABASE IF EXISTS "%s"`, snapshotName),
// Create a copy of the database to another database to use as a template now that it was fully migrated
fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, snapshotName, c.dbName, c.user),
// Snapshot the template database so we can restore it onto our original database going forward
fmt.Sprintf(`ALTER DATABASE "%s" WITH is_template = TRUE`, snapshotName),
}

for _, cmd := range cmds {
exitCode, reader, err := c.Exec(ctx, []string{"psql", "-U", c.user, "-d", c.dbName, "-c", cmd})
if err != nil {
return err
}
if exitCode != 0 {
buf := new(strings.Builder)
_, err := io.Copy(buf, reader)
if err != nil {
return fmt.Errorf("non-zero exit code for snapshot command, could not read command output: %w", err)
}

return fmt.Errorf("non-zero exit code for snapshot command: %s", buf.String())
}
); err != nil {
return err
}

c.snapshotName = snapshotName

return nil
}

// Restore will restore the database to a specific snapshot. By default, it will restore the last snapshot taken on the
// database by the Snapshot method. If a snapshot name is provided, it will instead try to restore the snapshot by name.
func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption) error {
snapshotName, err := c.checkSnapshotConfig(opts)
if err != nil {
return err
}

// execute the commands to restore the snapshot, in order
return c.execCommandsSQL(ctx,
// Drop the entire database by connecting to the postgres global database
fmt.Sprintf(`DROP DATABASE "%s" with (FORCE)`, c.dbName),
// Then restore the previous snapshot
fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, c.dbName, snapshotName, c.user),
)
}

func (c *PostgresContainer) checkSnapshotConfig(opts []SnapshotOption) (string, error) {
config := &snapshotConfig{}
for _, opt := range opts {
config = opt(config)
@@ -246,17 +238,61 @@ func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption)
}

if c.dbName == "postgres" {
return fmt.Errorf("cannot restore the postgres system database as it cannot be dropped to be restored")
return "", fmt.Errorf("cannot restore the postgres system database as it cannot be dropped to be restored")
}
return snapshotName, nil
}

// execute the commands to restore the snapshot, in order
cmds := []string{
// Drop the entire database by connecting to the postgres global database
fmt.Sprintf(`DROP DATABASE "%s" with (FORCE)`, c.dbName),
// Then restore the previous snapshot
fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, c.dbName, snapshotName, c.user),
func (c *PostgresContainer) execCommandsSQL(ctx context.Context, cmds ...string) error {
conn, cleanup, err := c.snapshotConnection(ctx)
if err != nil {
testcontainers.Logger.Printf("Could not connect to database to restore snapshot, falling back to `docker exec psql`: %v", err)
return c.execCommandsFallback(ctx, cmds)
}
defer cleanup()
mdelapenya marked this conversation as resolved.
Show resolved Hide resolved
for _, cmd := range cmds {
if _, err := conn.ExecContext(ctx, cmd); err != nil {
return fmt.Errorf("could not execute restore command %s: %w", cmd, err)
}
}
return nil
}

// snapshotConnection connects to the actual database using the "postgres" sql.DB driver, if it exists.
// The returned function should be called as a defer() to close the pool.
// No need to close the individual connection, that is done as part of the pool close.
// Also, no need to cache the connection pool, since it is a single connection which is very fast to establish.
func (c *PostgresContainer) snapshotConnection(ctx context.Context) (*sql.Conn, func(), error) {
// Connect to the database "postgres" instead of the app one
c2 := &PostgresContainer{
Container: c.Container,
dbName: "postgres",
user: c.user,
password: c.password,
}

// Try to use an actual postgres connection, if the driver is loaded
connStr := c2.MustConnectionString(ctx, "sslmode=disable")
pool, err := sql.Open("postgres", connStr)
if err != nil {
return nil, nil, fmt.Errorf("sql.Open for snapshot connection failed: %w", err)
}

cleanupPool := func() {
if err := pool.Close(); err != nil {
testcontainers.Logger.Printf("Could not close database connection pool after restoring snapshot: %v", err)
}
}

conn, err := pool.Conn(ctx)
if err != nil {
cleanupPool()
return nil, nil, fmt.Errorf("DB.Conn for snapshot connection failed: %w", err)
}
return conn, cleanupPool, nil
}

func (c *PostgresContainer) execCommandsFallback(ctx context.Context, cmds []string) error {
for _, cmd := range cmds {
exitCode, reader, err := c.Exec(ctx, []string{"psql", "-v", "ON_ERROR_STOP=1", "-U", c.user, "-d", "postgres", "-c", cmd})
if err != nil {
@@ -272,6 +308,5 @@ func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption)
return fmt.Errorf("non-zero exit code for restore command: %s", buf.String())
}
}

return nil
}