feature: custom ForSQL query

container_test.go

@@ -12,9 +12,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/docker/go-connections/nat"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/testcontainers/testcontainers-go/wait"
+
+	_ "github.com/lib/pq"
 )
 
 func Test_ContainerValidation(t *testing.T) {
@@ -414,3 +417,79 @@ func TestVolumeMount(t *testing.T) {
 		})
 	}
 }
+
+func TestContainerWithWaitForSQL(t *testing.T) {
+	const dbname = "postgres"
+
+	ctx := context.Background()
+
+	var env = map[string]string{
+		"POSTGRES_PASSWORD": "password",
+		"POSTGRES_USER":     "postgres",
+		"POSTGRES_DB":       dbname,
+	}
+	var port = "5432/tcp"
+	dbURL := func(port nat.Port) string {
+		return fmt.Sprintf("postgres://postgres:password@localhost:%s/%s?sslmode=disable", port.Port(), dbname)
+	}
+
+	t.Run("default query", func(t *testing.T) {
+		req := ContainerRequest{
+			Image:        "postgres:14.1-alpine",
+			ExposedPorts: []string{port},
+			Cmd:          []string{"postgres", "-c", "fsync=off"},
+			Env:          env,
+			WaitingFor: wait.ForSQL(nat.Port(port), "postgres", dbURL).
+				Timeout(time.Second * 5),
+		}
+		container, err := GenericContainer(ctx, GenericContainerRequest{
+			ContainerRequest: req,
+			Started:          true,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		defer container.Terminate(ctx)
+	})
+	t.Run("custom query", func(t *testing.T) {
+		req := ContainerRequest{
+			Image:        "postgres:14.1-alpine",
+			ExposedPorts: []string{port},
+			Cmd:          []string{"postgres", "-c", "fsync=off"},
+			Env:          env,
+			WaitingFor: wait.ForSQL(nat.Port(port), "postgres", dbURL).
+				Timeout(time.Second * 5).
+				WithQuery("SELECT 10"),
+		}
+		container, err := GenericContainer(ctx, GenericContainerRequest{
+			ContainerRequest: req,
+			Started:          true,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		defer container.Terminate(ctx)
+	})
+	t.Run("custom bad query", func(t *testing.T) {
+		req := ContainerRequest{
+			Image:        "postgres:14.1-alpine",
+			ExposedPorts: []string{port},
+			Cmd:          []string{"postgres", "-c", "fsync=off"},
+			Env:          env,
+			WaitingFor: wait.ForSQL(nat.Port(port), "postgres", dbURL).
+				Timeout(time.Second * 5).
+				WithQuery("SELECT 'a' from b"),
+		}
+		container, err := GenericContainer(ctx, GenericContainerRequest{
+			ContainerRequest: req,
+			Started:          true,
+		})
+		if err == nil {
+			t.Fatal("expected error, but got a nil")
+		}
+
+		defer container.Terminate(ctx)
+	})
+}

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/uuid v1.3.0
+	github.com/lib/pq v1.10.6
 	github.com/magiconair/properties v1.8.6
 	github.com/moby/sys/mount v0.2.0 // indirect
 	github.com/moby/sys/mountinfo v0.5.0 // indirect

go.sum

@@ -422,6 +422,8 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/pty v1.1.5/go.mod h1:9r2w37qlBe7rQ6e1fg1S/9xpWHSnaqNdHD3WcMdbPDA=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/lib/pq v1.10.6 h1:jbk+ZieJ0D7EVGJYpL9QTz7/YW6UHbmdnZWYyK5cdBs=
+github.com/lib/pq v1.10.6/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/magiconair/properties v1.8.6 h1:5ibWZ6iY0NctNGWo87LalDlEZ6R41TqbbDamhfG/Qzo=
 github.com/magiconair/properties v1.8.6/go.mod h1:y3VJvCyxH9uVvJTWEGAELF3aiYNyPKd5NZ3oSwXrF60=

wait/sql.go

@@ -9,6 +9,8 @@ import (
 	"github.com/docker/go-connections/nat"
 )
 
+const defaultForSqlQuery = "SELECT 1"
+
 //ForSQL constructs a new waitForSql strategy for the given driver
 func ForSQL(port nat.Port, driver string, url func(nat.Port) string) *waitForSql {
 	return &waitForSql{
@@ -17,6 +19,7 @@ func ForSQL(port nat.Port, driver string, url func(nat.Port) string) *waitForSql
 		Driver:         driver,
 		startupTimeout: defaultStartupTimeout(),
 		PollInterval:   defaultPollInterval(),
+		query:          defaultForSqlQuery,
 	}
 }
 
@@ -26,6 +29,7 @@ type waitForSql struct {
 	Port           nat.Port
 	startupTimeout time.Duration
 	PollInterval   time.Duration
+	query          string
 }
 
 //Timeout sets the maximum waiting time for the strategy after which it'll give up and return an error
@@ -40,8 +44,15 @@ func (w *waitForSql) WithPollInterval(pollInterval time.Duration) *waitForSql {
 	return w
 }
 
-//WaitUntilReady repeatedly tries to run "SELECT 1" query on the given port using sql and driver.
-// If the it doesn't succeed until the timeout value which defaults to 60 seconds, it will return an error
+//WithQuery can be used to override the default query used in the strategy.
+func (w *waitForSql) WithQuery(query string) *waitForSql {
+	w.query = query
+	return w
+}
+
+//WaitUntilReady repeatedly tries to run "SELECT 1" or user defined query on the given port using sql and driver.
+//
+// If it doesn't succeed until the timeout value which defaults to 60 seconds, it will return an error.
 func (w *waitForSql) WaitUntilReady(ctx context.Context, target StrategyTarget) (err error) {
 	ctx, cancel := context.WithTimeout(ctx, w.startupTimeout)
 	defer cancel()
@@ -72,7 +83,7 @@ func (w *waitForSql) WaitUntilReady(ctx context.Context, target StrategyTarget)
 			return ctx.Err()
 		case <-ticker.C:
 
-			if _, err := db.ExecContext(ctx, "SELECT 1"); err != nil {
+			if _, err := db.ExecContext(ctx, w.query); err != nil {
 				continue
 			}
 			return nil

wait/sql_test.go

@@ -0,0 +1,30 @@
+package wait
+
+import (
+	"testing"
+
+	"github.com/docker/go-connections/nat"
+)
+
+func Test_waitForSql_WithQuery(t *testing.T) {
+	t.Run("default query", func(t *testing.T) {
+		w := ForSQL("5432/tcp", "postgres", func(port nat.Port) string {
+			return "fake-url"
+		})
+
+		if got := w.query; got != defaultForSqlQuery {
+			t.Fatalf("expected %s, got %s", defaultForSqlQuery, got)
+		}
+	})
+	t.Run("custom query", func(t *testing.T) {
+		const q = "SELECT 100;"
+
+		w := ForSQL("5432/tcp", "postgres", func(port nat.Port) string {
+			return "fake-url"
+		}).WithQuery(q)
+
+		if got := w.query; got != q {
+			t.Fatalf("expected %s, got %s", q, got)
+		}
+	})
+}