scroll/database/docker/docker_db.go

package docker

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/client"

	"scroll-tech/common/docker"
)

var (
	cli *client.Client
)

func init() {
	var err error
	cli, err = client.NewClientWithOpts(client.FromEnv)
	if err != nil {
		panic(err)
	}
	cli.NegotiateAPIVersion(context.Background())
}

// ImgDB the postgres image manager.
type ImgDB struct {
	image string
	name  string
	id    string

	dbName   string
	port     int
	password string

	running bool
	*docker.Cmd
}

// NewImgDB return postgres db img instance.
func NewImgDB(t *testing.T, image, password, dbName string, port int) docker.ImgInstance {
	return &ImgDB{
		image:    image,
		name:     fmt.Sprintf("%s-%s_%d", image, dbName, port),
		password: password,
		dbName:   dbName,
		port:     port,
		Cmd:      docker.NewCmd(t),
	}
}

// Start postgres db container.
func (i *ImgDB) Start() error {
	id := docker.GetContainerID(i.name)
	if id != "" {
		return fmt.Errorf("container already exist, name: %s", i.name)
	}
	i.Cmd.RunCmd(i.prepare(), true)
	i.running = i.isOk()
	if !i.running {
		_ = i.Stop()
		return fmt.Errorf("failed to start image: %s", i.image)
	}
	return nil
}

// Stop the container.
func (i *ImgDB) Stop() error {
	if !i.running {
		return nil
	}
	i.running = false

	ctx := context.Background()
	// check if container is running, stop the running container.
	id := docker.GetContainerID(i.name)
	if id != "" {
		timeout := time.Second * 3
		if err := cli.ContainerStop(ctx, id, &timeout); err != nil {
			return err
		}
		i.id = id
	}
	// remove the stopped container.
	return cli.ContainerRemove(ctx, i.id, types.ContainerRemoveOptions{})
}

// Endpoint return the dsn.
func (i *ImgDB) Endpoint() string {
	if !i.running {
		return ""
	}
	return fmt.Sprintf("postgres://postgres:%s@localhost:%d/%s?sslmode=disable", i.password, i.port, i.dbName)
}

func (i *ImgDB) prepare() []string {
	cmd := []string{"docker", "run", "--name", i.name, "-p", fmt.Sprintf("%d:5432", i.port)}
	envs := []string{
		"-e", "POSTGRES_PASSWORD=" + i.password,
		"-e", fmt.Sprintf("POSTGRES_DB=%s", i.dbName),
	}

	cmd = append(cmd, envs...)
	return append(cmd, i.image)
}

func (i *ImgDB) isOk() bool {
	keyword := "database system is ready to accept connections"
	okCh := make(chan struct{}, 1)
	i.RegistFunc(keyword, func(buf string) {
		if strings.Contains(buf, keyword) {
			select {
			case okCh <- struct{}{}:
			default:
				return
			}
		}
	})
	defer i.UnRegistFunc(keyword)

	select {
	case <-okCh:
		time.Sleep(time.Millisecond * 1500)
		i.id = docker.GetContainerID(i.name)
		return i.id != ""
	case <-time.NewTimer(time.Second * 10).C:
		return false
	}
}