github.com/rudderlabs/rudder-go-kit@v0.30.0/testhelper/docker/resource/postgres/postgres.go (about)

     1  package postgres
     2  
     3  import (
     4  	"bytes"
     5  	"database/sql"
     6  	_ "encoding/json"
     7  	"fmt"
     8  
     9  	_ "github.com/lib/pq"
    10  	"github.com/ory/dockertest/v3"
    11  	"github.com/ory/dockertest/v3/docker"
    12  	dc "github.com/ory/dockertest/v3/docker"
    13  
    14  	"github.com/rudderlabs/rudder-go-kit/bytesize"
    15  	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource"
    16  )
    17  
    18  const (
    19  	postgresDefaultDB       = "jobsdb"
    20  	postgresDefaultUser     = "rudder"
    21  	postgresDefaultPassword = "password"
    22  )
    23  
    24  type Resource struct {
    25  	DB       *sql.DB
    26  	DBDsn    string
    27  	Database string
    28  	Password string
    29  	User     string
    30  	Host     string
    31  	Port     string
    32  
    33  	ContainerName string
    34  	ContainerID   string
    35  }
    36  
    37  func Setup(pool *dockertest.Pool, d resource.Cleaner, opts ...func(*Config)) (*Resource, error) {
    38  	c := &Config{
    39  		Tag:     "15-alpine",
    40  		ShmSize: 128 * bytesize.MB,
    41  	}
    42  	for _, opt := range opts {
    43  		opt(c)
    44  	}
    45  
    46  	cmd := []string{"postgres"}
    47  	for _, opt := range c.Options {
    48  		cmd = append(cmd, "-c", opt)
    49  	}
    50  	// pulls an image, creates a container based on it and runs it
    51  	postgresContainer, err := pool.RunWithOptions(&dockertest.RunOptions{
    52  		Repository: "postgres",
    53  		Tag:        c.Tag,
    54  		Env: []string{
    55  			"POSTGRES_PASSWORD=" + postgresDefaultPassword,
    56  			"POSTGRES_DB=" + postgresDefaultDB,
    57  			"POSTGRES_USER=" + postgresDefaultUser,
    58  		},
    59  		Cmd: cmd,
    60  	}, func(hc *dc.HostConfig) {
    61  		hc.ShmSize = c.ShmSize
    62  		hc.OOMKillDisable = c.OOMKillDisable
    63  		hc.Memory = c.Memory
    64  	})
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  	var db *sql.DB
    69  
    70  	d.Cleanup(func() {
    71  		if d.Failed() && c.PrintLogsOnError {
    72  			if c, found := pool.ContainerByName(postgresContainer.Container.Name); found {
    73  				d.Log(fmt.Sprintf("%q postgres container state: %+v", c.Container.Name, c.Container.State))
    74  				b := bytes.NewBufferString("")
    75  				if err := pool.Client.Logs(docker.LogsOptions{
    76  					Container:    c.Container.ID,
    77  					Stdout:       true,
    78  					Stderr:       true,
    79  					OutputStream: b,
    80  					ErrorStream:  b,
    81  				}); err != nil {
    82  					_, _ = b.Write([]byte(fmt.Sprintf("could not get logs: %s", err)))
    83  				}
    84  				d.Log(fmt.Sprintf("%q postgres container logs:\n%s", c.Container.Name, b.String()))
    85  			}
    86  		}
    87  		if err := pool.Purge(postgresContainer); err != nil {
    88  			d.Log("Could not purge resource:", err)
    89  		}
    90  		if db != nil {
    91  			_ = db.Close()
    92  		}
    93  	})
    94  
    95  	dbDSN := fmt.Sprintf(
    96  		"postgres://%s:%s@localhost:%s/%s?sslmode=disable",
    97  		postgresDefaultUser, postgresDefaultPassword, postgresContainer.GetPort("5432/tcp"), postgresDefaultDB,
    98  	)
    99  
   100  	// exponential backoff-retry, because the application in the container might not be ready to accept connections yet
   101  	err = pool.Retry(func() (err error) {
   102  		// 1. use pg_isready
   103  		var w bytes.Buffer
   104  		code, err := postgresContainer.Exec([]string{
   105  			"bash",
   106  			"-c",
   107  			fmt.Sprintf("pg_isready -d %[1]s -U %[2]s", postgresDefaultDB, postgresDefaultUser),
   108  		}, dockertest.ExecOptions{StdOut: &w, StdErr: &w})
   109  		if err != nil {
   110  			return err
   111  		}
   112  		if code != 0 {
   113  			return fmt.Errorf("postgres not ready:\n%s" + w.String())
   114  		}
   115  
   116  		// 2. create a sql.DB and verify connection
   117  		if db, err = sql.Open("postgres", dbDSN); err != nil {
   118  			return fmt.Errorf("opening database: %w", err)
   119  		}
   120  		defer func() {
   121  			if err != nil {
   122  				_ = db.Close()
   123  			}
   124  		}()
   125  		if err = db.Ping(); err != nil {
   126  			return fmt.Errorf("pinging database: %w", err)
   127  		}
   128  		var one int
   129  		if err = db.QueryRow("SELECT 1").Scan(&one); err != nil {
   130  			return fmt.Errorf("querying database: %w", err)
   131  		}
   132  		return nil
   133  	})
   134  	if err != nil {
   135  		return nil, fmt.Errorf("waiting for database to startup: %w", err)
   136  	}
   137  	return &Resource{
   138  		DB:            db,
   139  		DBDsn:         dbDSN,
   140  		Database:      postgresDefaultDB,
   141  		User:          postgresDefaultUser,
   142  		Password:      postgresDefaultPassword,
   143  		Host:          "localhost",
   144  		Port:          postgresContainer.GetPort("5432/tcp"),
   145  		ContainerName: postgresContainer.Container.Name,
   146  		ContainerID:   postgresContainer.Container.ID,
   147  	}, nil
   148  }