github.com/openfga/openfga@v1.5.4-rc1/pkg/testfixtures/storage/postgres.go (about)

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/cenkalti/backoff/v4"
    11  	"github.com/docker/docker/api/types/container"
    12  	_ "github.com/jackc/pgx/v5/stdlib" // PostgreSQL driver.
    13  	"github.com/pressly/goose/v3"
    14  	"github.com/stretchr/testify/require"
    15  	"github.com/testcontainers/testcontainers-go"
    16  	testcontainerspostgres "github.com/testcontainers/testcontainers-go/modules/postgres"
    17  	"github.com/testcontainers/testcontainers-go/wait"
    18  
    19  	"github.com/openfga/openfga/assets"
    20  )
    21  
    22  const (
    23  	postgresImage = "postgres:14"
    24  )
    25  
    26  type postgresTestContainer struct {
    27  	addr     string
    28  	version  int64
    29  	username string
    30  	password string
    31  }
    32  
    33  // NewPostgresTestContainer returns an implementation of the DatastoreTestContainer interface
    34  // for Postgres.
    35  func NewPostgresTestContainer() *postgresTestContainer {
    36  	return &postgresTestContainer{}
    37  }
    38  
    39  func (p *postgresTestContainer) GetDatabaseSchemaVersion() int64 {
    40  	return p.version
    41  }
    42  
    43  // RunPostgresTestContainer runs a Postgres container, connects to it, and returns a
    44  // bootstrapped implementation of the DatastoreTestContainer interface wired up for the
    45  // Postgres datastore engine.
    46  func (p *postgresTestContainer) RunPostgresTestContainer(t testing.TB) DatastoreTestContainer {
    47  	ctx := context.Background()
    48  
    49  	postgresContainer, err := testcontainerspostgres.RunContainer(ctx,
    50  		testcontainers.WithImage(postgresImage),
    51  		testcontainers.WithWaitStrategy(wait.
    52  			ForLog("database system is ready to accept connections").
    53  			WithOccurrence(2).
    54  			WithStartupTimeout(5*time.Second),
    55  		),
    56  		testcontainers.WithHostConfigModifier(func(hostConfig *container.HostConfig) {
    57  			hostConfig.Tmpfs = map[string]string{"/var/lib/postgresql/data": ""}
    58  		}),
    59  		testcontainerspostgres.WithDatabase("defaultdb"),
    60  		testcontainerspostgres.WithUsername("postgres"),
    61  		testcontainerspostgres.WithPassword("secret"),
    62  	)
    63  	require.NoError(t, err)
    64  	t.Cleanup(func() { require.NoError(t, postgresContainer.Terminate(ctx)) })
    65  
    66  	postgresHost, err := postgresContainer.Host(ctx)
    67  	require.NoError(t, err)
    68  	postgresPort, err := postgresContainer.MappedPort(ctx, "5432/tcp")
    69  	require.NoError(t, err)
    70  
    71  	pgTestContainer := &postgresTestContainer{
    72  		addr:     net.JoinHostPort(postgresHost, postgresPort.Port()),
    73  		username: "postgres",
    74  		password: "secret",
    75  	}
    76  
    77  	uri := fmt.Sprintf("postgres://%s:%s@%s/defaultdb?sslmode=disable", pgTestContainer.username, pgTestContainer.password, pgTestContainer.addr)
    78  
    79  	goose.SetLogger(goose.NopLogger())
    80  
    81  	db, err := goose.OpenDBWithDriver("pgx", uri)
    82  	require.NoError(t, err)
    83  	defer db.Close()
    84  
    85  	backoffPolicy := backoff.NewExponentialBackOff()
    86  	backoffPolicy.MaxElapsedTime = 30 * time.Second
    87  	err = backoff.Retry(
    88  		func() error {
    89  			return db.Ping()
    90  		},
    91  		backoffPolicy,
    92  	)
    93  	require.NoError(t, err, "failed to connect to postgres container")
    94  
    95  	goose.SetBaseFS(assets.EmbedMigrations)
    96  
    97  	err = goose.Up(db, assets.PostgresMigrationDir)
    98  	require.NoError(t, err)
    99  
   100  	version, err := goose.GetDBVersion(db)
   101  	require.NoError(t, err)
   102  	pgTestContainer.version = version
   103  
   104  	return pgTestContainer
   105  }
   106  
   107  // GetConnectionURI returns the postgres connection uri for the running postgres test container.
   108  func (p *postgresTestContainer) GetConnectionURI(includeCredentials bool) string {
   109  	creds := ""
   110  	if includeCredentials {
   111  		creds = fmt.Sprintf("%s:%s@", p.username, p.password)
   112  	}
   113  
   114  	return fmt.Sprintf(
   115  		"postgres://%s%s/%s?sslmode=disable",
   116  		creds,
   117  		p.addr,
   118  		"defaultdb",
   119  	)
   120  }
   121  
   122  func (p *postgresTestContainer) GetUsername() string {
   123  	return p.username
   124  }
   125  
   126  func (p *postgresTestContainer) GetPassword() string {
   127  	return p.password
   128  }