Repository: ariga/atlas
Pull request #676 (2022-03-31 08:22:01)

cmd/action: add MariaDB and PostgreSQL support for docker driver


Patch:
diff --git a/cmd/action/docker.go b/cmd/action/docker.go
index da17202e..b4ab5ebc 100644
--- a/cmd/action/docker.go
+++ b/cmd/action/docker.go
@@ -28,6 +28,7 @@ const pass = "pass"
 type (
 	// DockerConfig is used to configure container creation.
 	DockerConfig struct {
+		setup []string // contains statements to execute once the service is up
 		// Image is the name of the image to pull and run.
 		Image string
 		// Env vars to pass to the docker container.
@@ -65,14 +66,35 @@ func NewConfig(opts ...DockerConfigOption) (*DockerConfig, error) {
 
 // MySQL returns a new DockerConfig for a MySQL image.
 func MySQL(version string, opts ...DockerConfigOption) (*DockerConfig, error) {
-	return NewConfig(append(
-		[]DockerConfigOption{
-			Image("mysql:" + version),
-			Port("3306"),
-			Env("MYSQL_ROOT_PASSWORD=" + pass),
-		},
-		opts...,
-	)...,
+	return NewConfig(
+		append(
+			[]DockerConfigOption{
+				Image("mysql:" + version),
+				Port("3306"),
+				Env("MYSQL_ROOT_PASSWORD=" + pass),
+			},
+			opts...,
+		)...,
+	)
+}
+
+// MariaDB returns a new DockerConfig for a MariaDB image.
+func MariaDB(version string, opts ...DockerConfigOption) (*DockerConfig, error) {
+	return MySQL(version, append([]DockerConfigOption{Image("mariadb:" + version)}, opts...)...)
+}
+
+// PostgreSQL returns a new DockerConfig for a PostgreSQL image.
+func PostgreSQL(version string, opts ...DockerConfigOption) (*DockerConfig, error) {
+	return NewConfig(
+		append(
+			[]DockerConfigOption{
+				Image("postgres:" + version),
+				Port("5432"),
+				Env("POSTGRES_PASSWORD=" + pass),
+				setup("DROP SCHEMA IF EXISTS public CASCADE;"),
+			},
+			opts...,
+		)...,
 	)
 }
 
@@ -83,7 +105,7 @@ func MySQL(version string, opts ...DockerConfigOption) (*DockerConfig, error) {
 //
 func Image(i string) DockerConfigOption {
 	return func(c *DockerConfig) error {
-		c.Image = i
+		c.Image = strings.TrimSuffix(i, ":")
 		return nil
 	}
 }
@@ -124,6 +146,17 @@ func Out(w io.Writer) DockerConfigOption {
 	}
 }
 
+// setup adds statements to execute once the service is ready. For example:
+//
+//  setup("DROP SCHEMA IF EXISTS public CASCADE;")
+//
+func setup(s ...string) DockerConfigOption {
+	return func(c *DockerConfig) error {
+		c.setup = s
+		return nil
+	}
+}
+
 // Run pulls and starts a new docker container from the DockerConfig.
 func (c *DockerConfig) Run(ctx context.Context) (*Container, error) {
 	// Make sure the configuration is not missing critical values.
@@ -163,8 +196,8 @@ func (c *DockerConfig) Run(ctx context.Context) (*Container, error) {
 }
 
 // Down stops and removes this container.
-func (c *Container) Down(ctx context.Context) error {
-	return exec.CommandContext(ctx, "docker", "stop", c.ID).Run() //nolint:gosec
+func (c *Container) Down() error {
+	return exec.Command("docker", "stop", c.ID).Run() //nolint:gosec
 }
 
 // Wait waits for this container to be ready.
@@ -188,11 +221,18 @@ func (c *Container) Wait(ctx context.Context, timeout time.Duration) error {
 		case <-time.After(100 * time.Millisecond):
 			db, err := sql.Open(drv, dsn)
 			if err != nil {
-				continue
+				return err
 			}
 			if err := db.PingContext(ctx); err != nil {
 				continue
 			}
+			for _, s := range c.cfg.setup {
+				if _, err := db.ExecContext(ctx, s); err != nil {
+					_ = db.Close()
+					return fmt.Errorf("%q: %w", s, err)
+				}
+			}
+			_ = db.Close()
 			fmt.Fprintln(c.out, "Service is ready to connect!")
 			return nil
 		case <-ctx.Done():
@@ -205,16 +245,37 @@ func (c *Container) Wait(ctx context.Context, timeout time.Duration) error {
 
 // Driver returns the database/sql driver name.
 func (c *Container) Driver() string {
-	return strings.SplitN(c.cfg.Image, ":", 2)[0]
+	switch img := strings.SplitN(c.cfg.Image, ":", 2)[0]; img {
+	case "mysql", "mariadb":
+		return "mysql"
+	default:
+		return img
+	}
 }
 
 // DSN returns a DSN to connect to the Container.
 func (c *Container) DSN() (string, error) {
-	switch drv := c.Driver(); drv {
+	switch img := strings.SplitN(c.cfg.Image, ":", 2)[0]; img {
 	case "mysql", "mariadb":
 		return fmt.Sprintf("root:%s@tcp(localhost:%s)/", c.Passphrase, c.Port), nil
+	case "postgres":
+		return fmt.Sprintf("host=localhost port=%s dbname=postgres user=postgres password=%s sslmode=disable", c.Port, c.Passphrase), nil
 	default:
-		return "", fmt.Errorf("unknown driver: %q", drv)
+		return "", fmt.Errorf("unsupported image: %q", c.cfg.Image)
+	}
+}
+
+// URL returns a URL to connect to the Container.
+func (c *Container) URL() (string, error) {
+	switch img := strings.SplitN(c.cfg.Image, ":", 2)[0]; img {
+	case "postgres":
+		return fmt.Sprintf("postgres://postgres:%s@localhost:%s/postgres?sslmode=disable", c.Passphrase, c.Port), nil
+	default:
+		dsn, err := c.DSN()
+		if err != nil {
+			return "", err
+		}
+		return fmt.Sprintf("%s://%s", strings.SplitN(c.cfg.Image, ":", 2)[0], dsn), nil
 	}
 }
 
diff --git a/cmd/action/provider.go b/cmd/action/provider.go
index 0212fa6b..7227f8a0 100644
--- a/cmd/action/provider.go
+++ b/cmd/action/provider.go
@@ -82,16 +82,15 @@ func sqliteProvider(_ context.Context, dsn string, _ ...ProviderOption) (*Driver
 	}, nil
 }
 
-var reDockerConfig = regexp.MustCompile("(mysql|mariadb)(?::([0-9a-zA-Z.-]+)?(\\?.*)?)?")
+var reDockerConfig = regexp.MustCompile("(mysql|mariadb|tidb|postgres)(?::([0-9a-zA-Z.-]+)?(\\?.*)?)?")
 
 type dockerCloser struct {
 	c   *Container
 	drv *Driver
-	ctx context.Context
 }
 
 func (dc *dockerCloser) Close() (err error) {
-	derr, cerr := dc.drv.Close(), dc.c.Down(dc.ctx)
+	derr, cerr := dc.drv.Close(), dc.c.Down()
 	if derr != nil {
 		err = derr
 	}
@@ -105,14 +104,23 @@ func dockerProvider(ctx context.Context, dsn string, opts ...ProviderOption) (*D
 	// The DSN has the driver part (docker:// removed already.
 	// Get rid of the query arguments, and we have the image name.
 	m := reDockerConfig.FindStringSubmatch(dsn)
+	if len(m) != 4 {
+		return nil, fmt.Errorf("invalid docker url %q", dsn)
+	}
 	img, v := m[1], m[2]
 	var (
 		cfg *DockerConfig
 		err error
 	)
 	switch img {
-	case "mysql", "mariadb":
+	case "mysql":
 		cfg, err = MySQL(v)
+	case "mariadb":
+		cfg, err = MariaDB(v)
+	case "postgres":
+		cfg, err = PostgreSQL(v)
+	default:
+		return nil, fmt.Errorf("unsupported docker image %q", img)
 	}
 	if err != nil {
 		return nil, err
@@ -131,22 +139,23 @@ func dockerProvider(ctx context.Context, dsn string, opts ...ProviderOption) (*D
 		return nil, err
 	}
 	if err := c.Wait(ctx, time.Minute); err != nil {
+		_ = c.Down()
 		return nil, err
 	}
-	d, err := c.DSN()
+	u, err := c.URL()
 	if err != nil {
-		_ = c.Down(ctx)
+		_ = c.Down()
 		return nil, err
 	}
-	drv, err := DefaultMux.OpenAtlas(ctx, fmt.Sprintf("%s://%s", img, d))
+	drv, err := DefaultMux.OpenAtlas(ctx, u)
 	if err != nil {
-		_ = c.Down(ctx)
+		_ = c.Down()
 		return nil, err
 	}
 	return &Driver{
 		Driver:      drv.Driver,
 		Marshaler:   drv.Marshaler,
 		Unmarshaler: drv.Unmarshaler,
-		Closer:      &dockerCloser{c, drv, ctx},
+		Closer:      &dockerCloser{c, drv},
 	}, nil
 }


Test patch:
diff --git a/cmd/action/docker_test.go b/cmd/action/docker_test.go
index faa76d89..bc0f22f3 100644
--- a/cmd/action/docker_test.go
+++ b/cmd/action/docker_test.go
@@ -2,14 +2,13 @@
 // This source code is licensed under the Apache 2.0 license found
 // in the LICENSE file in the root directory of this source tree.
 
-package action_test
+package action
 
 import (
 	"context"
 	"io/ioutil"
 	"testing"
 
-	"ariga.io/atlas/cmd/action"
 	"github.com/stretchr/testify/require"
 )
 
@@ -17,16 +16,37 @@ func TestDockerConfig(t *testing.T) {
 	ctx := context.Background()
 
 	// invalid config
-	_, err := (&action.DockerConfig{}).Run(ctx)
+	_, err := (&DockerConfig{}).Run(ctx)
 	require.Error(t, err)
 
 	// MySQL
-	cfg, err := action.MySQL("latest", action.Out(ioutil.Discard))
+	cfg, err := MySQL("latest", Out(ioutil.Discard))
 	require.NoError(t, err)
-	require.Equal(t, &action.DockerConfig{
+	require.Equal(t, &DockerConfig{
 		Image: "mysql:latest",
 		Env:   []string{"MYSQL_ROOT_PASSWORD=pass"},
 		Port:  "3306",
 		Out:   ioutil.Discard,
 	}, cfg)
+
+	// MariaDB
+	cfg, err = MariaDB("latest", Out(ioutil.Discard))
+	require.NoError(t, err)
+	require.Equal(t, &DockerConfig{
+		Image: "mariadb:latest",
+		Env:   []string{"MYSQL_ROOT_PASSWORD=pass"},
+		Port:  "3306",
+		Out:   ioutil.Discard,
+	}, cfg)
+
+	// PostgreSQL
+	cfg, err = PostgreSQL("latest", Out(ioutil.Discard))
+	require.NoError(t, err)
+	require.Equal(t, &DockerConfig{
+		Image: "postgres:latest",
+		Env:   []string{"POSTGRES_PASSWORD=pass"},
+		Port:  "5432",
+		Out:   ioutil.Discard,
+		setup: []string{"DROP SCHEMA IF EXISTS public CASCADE;"},
+	}, cfg)
 }