Repository: xataio/pgstream
Pull request #144 (2025-02-14 13:18:53)

Add pg_dump/pg_restore to internal library
This PR adds the `pg_dump`/`pg_restore` functionalities to the internal postgres library. A testcontainers package is created in order to reuse the setup for integration tests.


Patch:
diff --git a/internal/postgres/pg_dump.go b/internal/postgres/pg_dump.go
new file mode 100644
index 0000000..29e326e
--- /dev/null
+++ b/internal/postgres/pg_dump.go
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"bytes"
+	"fmt"
+	"os/exec"
+)
+
+var pgDumpCmd = "pg_dump"
+
+type PGDumpOptions struct {
+	// ConnectionString
+	ConnectionString string
+	// Format (c custom, d directory, t tar, p plain text)
+	Format string
+	// Schemas to export
+	Schemas []string
+	// Tables to export
+	Tables []string
+	// Tables to be excluded from the dump, regex patterns supported
+	ExcludeTables []string
+	// SchemaOnly if true, only schema will be exported (no data)
+	SchemaOnly bool
+	// do not dump privileges (grant/revoke)
+	NoPrivileges bool
+	// Options to pass to pg_dump
+	Options []string
+}
+
+func (opts *PGDumpOptions) ToArgs() []string {
+	options := make([]string, 0, len(opts.Schemas)+len(opts.Tables)+len(opts.ExcludeTables)+len(opts.Options))
+
+	options = append(options, opts.ConnectionString)
+
+	if opts.Format != "" {
+		options = append(options, fmt.Sprintf(`-F%v`, opts.Format))
+	}
+
+	if opts.SchemaOnly {
+		options = append(options, "--schema-only")
+	}
+
+	if opts.NoPrivileges {
+		options = append(options, "--no-privileges")
+	}
+
+	for _, schema := range opts.Schemas {
+		options = append(options, fmt.Sprintf("--schema=%v", schema))
+	}
+
+	for _, table := range opts.Tables {
+		options = append(options, fmt.Sprintf("--table=%v", table))
+	}
+
+	for _, table := range opts.ExcludeTables {
+		options = append(options, fmt.Sprintf("--exclude-table=%v", table))
+	}
+
+	options = append(options, opts.Options...)
+	return options
+}
+
+// Func RunPGDump runs pg_dump command with the given options and returns the result.
+func RunPGDump(opts PGDumpOptions) ([]byte, error) {
+	cmd := exec.Command(pgDumpCmd, opts.ToArgs()...) //nolint:gosec
+
+	// TODO: add streaming support when it needs to be used for large data dumps
+	var stderr bytes.Buffer
+	cmd.Stderr = &stderr
+
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, fmt.Errorf("error running pg_dump: %w: %s", err, stderr.String())
+	}
+
+	return out, nil
+}
diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go
new file mode 100644
index 0000000..efc6bea
--- /dev/null
+++ b/internal/postgres/pg_restore.go
@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os/exec"
+)
+
+var pgRestoreCmd = "pg_restore"
+
+type PGRestoreOptions struct {
+	// ConnectionString
+	ConnectionString string
+	// SchemaOnly if true, only schema will be restored (no data)
+	SchemaOnly bool
+	// Clean all the objects that will be restored
+	Clean bool
+	// Options to pass to pg_restore
+	Options []string
+}
+
+func (opts PGRestoreOptions) ToArgs() []string {
+	var options []string
+
+	options = append(options, "-d", opts.ConnectionString)
+
+	if opts.SchemaOnly {
+		options = append(options, "--schema-only")
+	}
+
+	if opts.Clean {
+		options = append(options, "--clean")
+		options = append(options, "--if-exists")
+	}
+
+	options = append(options, opts.Options...)
+	return options
+}
+
+// Func RunPGRestore runs pg_restore command with the given options and returns the result.
+func RunPGRestore(opts PGRestoreOptions, dump []byte) (string, error) {
+	cmd := exec.Command(pgRestoreCmd, opts.ToArgs()...) //nolint:gosec
+
+	var stderr bytes.Buffer
+	cmd.Stderr = &stderr
+
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return "", fmt.Errorf("error getting stdin pipe: %w", err)
+	}
+
+	go func() {
+		defer stdin.Close()
+		io.Copy(stdin, bytes.NewReader(dump))
+	}()
+
+	// TODO: add streaming support when large data output is required
+	out, err := cmd.Output()
+	if err != nil {
+		return "", fmt.Errorf("error running pg_restore: %w: (%s)", err, stderr.String())
+	}
+
+	return string(out), nil
+}


Test patch:
diff --git a/internal/postgres/pg_dump_pg_restore_integration_test.go b/internal/postgres/pg_dump_pg_restore_integration_test.go
new file mode 100644
index 0000000..edc83ee
--- /dev/null
+++ b/internal/postgres/pg_dump_pg_restore_integration_test.go
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/internal/testcontainers"
+)
+
+func Test_pgdump_pgrestore(t *testing.T) {
+	if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" {
+		t.Skip("skipping integration test...")
+	}
+
+	ctx := context.Background()
+
+	var sourcePGURL, targetPGURL string
+	cleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourcePGURL, testcontainers.Postgres17)
+	require.NoError(t, err)
+	defer cleanup()
+
+	cleanup2, err := testcontainers.SetupPostgresContainer(ctx, &targetPGURL, testcontainers.Postgres17)
+	require.NoError(t, err)
+	defer cleanup2()
+
+	sourceConn, err := NewConn(ctx, sourcePGURL)
+	require.NoError(t, err)
+	targetConn, err := NewConn(ctx, targetPGURL)
+	require.NoError(t, err)
+
+	testSchema := "test_schema"
+	testTable1 := "test_table_1"
+	testTable2 := "test_table_2"
+
+	// create schema + tables
+	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create schema %s", testSchema))
+	require.NoError(t, err)
+	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable1))
+	require.NoError(t, err)
+	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable2))
+	require.NoError(t, err)
+	// insert data
+	_, err = sourceConn.Exec(ctx, fmt.Sprintf("insert into %s.%s(name) values('a'),('b'),('c')", testSchema, testTable1))
+	require.NoError(t, err)
+
+	pgdumpOpts := PGDumpOptions{
+		ConnectionString: sourcePGURL,
+		Format:           "c",
+		SchemaOnly:       true,
+		Schemas:          []string{testSchema},
+		ExcludeTables:    []string{testSchema + "." + testTable2},
+	}
+
+	dump, err := RunPGDump(pgdumpOpts)
+	require.NoError(t, err)
+
+	pgrestoreOpts := PGRestoreOptions{
+		ConnectionString: targetPGURL,
+		SchemaOnly:       true,
+	}
+
+	_, err = RunPGRestore(pgrestoreOpts, dump)
+	require.NoError(t, err)
+
+	// schema only pgdump, no data should be available but the schema and
+	// selected table should exist.
+	var count int
+	err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable1)).Scan(&count)
+	require.NoError(t, err)
+	require.Equal(t, 0, count)
+	// test table 2 should not exist
+	err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable2)).Scan(&count)
+	require.Error(t, err)
+}
diff --git a/internal/testcontainers/test_elasticsearch_container.go b/internal/testcontainers/test_elasticsearch_container.go
new file mode 100644
index 0000000..9fc2270
--- /dev/null
+++ b/internal/testcontainers/test_elasticsearch_container.go
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package testcontainers
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/elasticsearch"
+)
+
+func SetupElasticsearchContainer(ctx context.Context, url *string) (cleanup, error) {
+	ctr, err := elasticsearch.Run(ctx, "docker.elastic.co/elasticsearch/elasticsearch:8.9.0",
+		testcontainers.WithEnv(map[string]string{"xpack.security.enabled": "false"})) // disable TLS
+	if err != nil {
+		return nil, fmt.Errorf("failed to start elasticsearch container: %w", err)
+	}
+
+	*url = ctr.Settings.Address
+
+	return func() error {
+		return ctr.Terminate(ctx)
+	}, nil
+}
diff --git a/internal/testcontainers/test_kafka_container.go b/internal/testcontainers/test_kafka_container.go
new file mode 100644
index 0000000..cfc2419
--- /dev/null
+++ b/internal/testcontainers/test_kafka_container.go
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package testcontainers
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/kafka"
+	"github.com/testcontainers/testcontainers-go/wait"
+)
+
+func SetupKafkaContainer(ctx context.Context, brokers *[]string) (cleanup, error) {
+	ctr, err := kafka.RunContainer(ctx,
+		kafka.WithClusterID("test-cluster"),
+		testcontainers.WithImage("confluentinc/confluent-local:7.5.0"),
+		testcontainers.WithWaitStrategy(
+			wait.ForLog("Kafka Server started").
+				WithOccurrence(1).
+				WithStartupTimeout(5*time.Second),
+		),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start kafka container: %w", err)
+	}
+
+	*brokers, err = ctr.Brokers(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving brokers for kafka container: %w", err)
+	}
+
+	return func() error {
+		return ctr.Terminate(ctx)
+	}, nil
+}
diff --git a/internal/testcontainers/test_opensearch_container.go b/internal/testcontainers/test_opensearch_container.go
new file mode 100644
index 0000000..5265e2a
--- /dev/null
+++ b/internal/testcontainers/test_opensearch_container.go
@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package testcontainers
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/opensearch"
+)
+
+func SetupOpenSearchContainer(ctx context.Context, url *string) (cleanup, error) {
+	ctr, err := opensearch.RunContainer(ctx,
+		testcontainers.WithImage("opensearchproject/opensearch:2.11.1"),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start opensearch container: %w", err)
+	}
+
+	*url, err = ctr.Address(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving url for opensearch container: %w", err)
+	}
+
+	return func() error {
+		return ctr.Terminate(ctx)
+	}, nil
+}
diff --git a/internal/testcontainers/test_postgres_container.go b/internal/testcontainers/test_postgres_container.go
new file mode 100644
index 0000000..d43a8e9
--- /dev/null
+++ b/internal/testcontainers/test_postgres_container.go
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package testcontainers
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
+	"github.com/testcontainers/testcontainers-go/wait"
+)
+
+type cleanup func() error
+
+type PostgresImage string
+
+const (
+	Postgres14 PostgresImage = "debezium/postgres:14-alpine"
+	Postgres17 PostgresImage = "debezium/postgres:17-alpine"
+)
+
+func SetupPostgresContainer(ctx context.Context, url *string, image PostgresImage, configFile ...string) (cleanup, error) {
+	waitForLogs := wait.
+		ForLog("database system is ready to accept connections").
+		WithOccurrence(2).
+		WithStartupTimeout(5 * time.Second)
+
+	opts := []testcontainers.ContainerCustomizer{
+		testcontainers.WithImage(string(image)),
+		testcontainers.WithWaitStrategy(waitForLogs),
+	}
+
+	if len(configFile) > 0 {
+		opts = append(opts, postgres.WithConfigFile(configFile[0]))
+	}
+	ctr, err := postgres.RunContainer(ctx, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start postgres container: %w", err)
+	}
+
+	*url, err = ctr.ConnectionString(ctx, "sslmode=disable")
+	if err != nil {
+		return nil, fmt.Errorf("retrieving connection string for postgres container: %w", err)
+	}
+
+	return func() error {
+		return ctr.Terminate(ctx)
+	}, nil
+}
diff --git a/pkg/stream/integration/setup_test.go b/pkg/stream/integration/setup_test.go
index b7625cb..2d0cf41 100644
--- a/pkg/stream/integration/setup_test.go
+++ b/pkg/stream/integration/setup_test.go
@@ -4,19 +4,11 @@ package integration
 
 import (
 	"context"
-	"fmt"
 	"log"
 	"os"
 	"testing"
-	"time"
-
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/modules/elasticsearch"
-	"github.com/testcontainers/testcontainers-go/modules/kafka"
-	"github.com/testcontainers/testcontainers-go/modules/opensearch"
-	"github.com/testcontainers/testcontainers-go/modules/postgres"
-	"github.com/testcontainers/testcontainers-go/wait"
 
+	"github.com/xataio/pgstream/internal/testcontainers"
 	"github.com/xataio/pgstream/pkg/stream"
 )
 
@@ -24,7 +16,7 @@ func TestMain(m *testing.M) {
 	// if integration tests are not enabled, nothing to setup
 	if os.Getenv("PGSTREAM_INTEGRATION_TESTS") != "" {
 		ctx := context.Background()
-		pgcleanup, err := setupPostgresContainer(ctx, &pgurl)
+		pgcleanup, err := testcontainers.SetupPostgresContainer(ctx, &pgurl, testcontainers.Postgres14, "config/postgresql.conf")
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -34,25 +26,25 @@ func TestMain(m *testing.M) {
 			log.Fatal(err)
 		}
 
-		kafkacleanup, err := setupKafkaContainer(ctx)
+		kafkacleanup, err := testcontainers.SetupKafkaContainer(ctx, &kafkaBrokers)
 		if err != nil {
 			log.Fatal(err)
 		}
 		defer kafkacleanup()
 
-		oscleanup, err := setupOpenSearchContainer(ctx)
+		oscleanup, err := testcontainers.SetupOpenSearchContainer(ctx, &opensearchURL)
 		if err != nil {
 			log.Fatal(err)
 		}
 		defer oscleanup()
 
-		escleanup, err := setupElasticsearchContainer(ctx)
+		escleanup, err := testcontainers.SetupElasticsearchContainer(ctx, &elasticsearchURL)
 		if err != nil {
 			log.Fatal(err)
 		}
 		defer escleanup()
 
-		targetPGCleanup, err := setupPostgresContainer(ctx, &targetPGURL)
+		targetPGCleanup, err := testcontainers.SetupPostgresContainer(ctx, &targetPGURL, testcontainers.Postgres14, "config/postgresql.conf")
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -61,86 +53,3 @@ func TestMain(m *testing.M) {
 
 	os.Exit(m.Run())
 }
-
-type cleanup func() error
-
-func setupPostgresContainer(ctx context.Context, url *string) (cleanup, error) {
-	waitForLogs := wait.
-		ForLog("database system is ready to accept connections").
-		WithOccurrence(2).
-		WithStartupTimeout(5 * time.Second)
-
-	ctr, err := postgres.RunContainer(ctx,
-		testcontainers.WithImage("debezium/postgres:14-alpine"),
-		testcontainers.WithWaitStrategy(waitForLogs),
-		postgres.WithConfigFile("config/postgresql.conf"),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("failed to start postgres container: %w", err)
-	}
-
-	*url, err = ctr.ConnectionString(ctx, "sslmode=disable")
-	if err != nil {
-		return nil, fmt.Errorf("retrieving connection string for postgres container: %w", err)
-	}
-
-	return func() error {
-		return ctr.Terminate(ctx)
-	}, nil
-}
-
-func setupKafkaContainer(ctx context.Context) (cleanup, error) {
-	ctr, err := kafka.RunContainer(ctx,
-		kafka.WithClusterID("test-cluster"),
-		testcontainers.WithImage("confluentinc/confluent-local:7.5.0"),
-		testcontainers.WithWaitStrategy(
-			wait.ForLog("Kafka Server started").
-				WithOccurrence(1).
-				WithStartupTimeout(5*time.Second),
-		),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("failed to start kafka container: %w", err)
-	}
-
-	kafkaBrokers, err = ctr.Brokers(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("retrieving brokers for kafka container: %w", err)
-	}
-
-	return func() error {
-		return ctr.Terminate(ctx)
-	}, nil
-}
-
-func setupOpenSearchContainer(ctx context.Context) (cleanup, error) {
-	ctr, err := opensearch.RunContainer(ctx,
-		testcontainers.WithImage("opensearchproject/opensearch:2.11.1"),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("failed to start opensearch container: %w", err)
-	}
-
-	opensearchURL, err = ctr.Address(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("retrieving url for opensearch container: %w", err)
-	}
-
-	return func() error {
-		return ctr.Terminate(ctx)
-	}, nil
-}
-
-func setupElasticsearchContainer(ctx context.Context) (cleanup, error) {
-	ctr, err := elasticsearch.Run(ctx, "docker.elastic.co/elasticsearch/elasticsearch:8.9.0",
-		testcontainers.WithEnv(map[string]string{"xpack.security.enabled": "false"})) // disable TLS
-	if err != nil {
-		return nil, fmt.Errorf("failed to start elasticsearch container: %w", err)
-	}
-
-	elasticsearchURL = ctr.Settings.Address
-
-	return func() error {
-		return ctr.Terminate(ctx)
-	}, nil
-}