From d198f78f1d1f18933750826f3f07f209c87b4a0b Mon Sep 17 00:00:00 2001
From: Brandon Sprague <brandon00sprague@gmail.com>
Date: Sun, 4 Dec 2022 20:00:36 -0800
Subject: [PATCH] Fix Postgres Docker flake and update 14.4 -> 14.6 (#4)

This PR fixes an issue with our Docker Postgres connection where we sometimes connect before Postgres is fully up, see [1] for more details.

[1] https://github.com/docker-library/postgres/issues/146
---
 example/golden/raw_dump.sql |  4 ++--
 testpgx.go                  | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/example/golden/raw_dump.sql b/example/golden/raw_dump.sql
index b885106..966dee2 100644
--- a/example/golden/raw_dump.sql
+++ b/example/golden/raw_dump.sql
@@ -2,8 +2,8 @@
 -- PostgreSQL database dump
 --
 
--- Dumped from database version 14.4 (Debian 14.4-1.pgdg110+1)
--- Dumped by pg_dump version 14.4 (Debian 14.4-1.pgdg110+1)
+-- Dumped from database version 14.6 (Debian 14.6-1.pgdg110+1)
+-- Dumped by pg_dump version 14.6 (Debian 14.6-1.pgdg110+1)
 
 SET statement_timeout = 0;
 SET lock_timeout = 0;
diff --git a/testpgx.go b/testpgx.go
index 9ef7e5f..16bfd14 100644
--- a/testpgx.go
+++ b/testpgx.go
@@ -1,6 +1,7 @@
 package testpgx
 
 import (
+	"bufio"
 	"context"
 	"database/sql"
 	"errors"
@@ -19,8 +20,8 @@ import (
 	"github.com/google/uuid"
 	"github.com/hashicorp/go-multierror"
 	"github.com/jackc/pgx/v4"
-	"github.com/jackc/pgx/v4/pgxpool"
 
+	"github.com/jackc/pgx/v4/pgxpool"
 	pgxstdlib "github.com/jackc/pgx/v4/stdlib"
 )
 
@@ -78,7 +79,7 @@ const (
 	testDBUser = "postgres"
 	testDBPass = "anypassword"
 
-	defaultPostgresImage = "postgres:14.4"
+	defaultPostgresImage = "postgres:14.6"
 	defaultMaxDBs        = 10
 )
 
@@ -180,6 +181,31 @@ func New(ctx context.Context, opts ...Option) (*Env, error) {
 		return nil, fmt.Errorf("when getting postgres cid: %w", err)
 	}
 
+	// Postgres starts up twice, first as an initialization phase, and then for
+	// real. See this thread [1] for more info.
+	// [1] https://github.com/docker-library/postgres/issues/146
+	pgLogCtx, pgLogDone := context.WithCancel(ctx)
+	defer pgLogDone()
+	pgLogCmd := exec.CommandContext(pgLogCtx, o.dockerBinaryPath, "logs", "-f", cID)
+	pgLogs, err := pgLogCmd.StdoutPipe()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get stdout pipe for docker logs: %w", err)
+	}
+	if err := pgLogCmd.Start(); err != nil {
+		return nil, fmt.Errorf("failed to start docker logs: %w", err)
+	}
+	initDone := make(chan struct{})
+	go func() {
+		sc := bufio.NewScanner(pgLogs)
+		for sc.Scan() {
+			if strings.Contains(sc.Text(), "PostgreSQL init process complete; ready for start up.") {
+				break
+			}
+		}
+		close(initDone)
+		pgLogDone()
+	}()
+
 	env := &Env{
 		postgresCid: cID,
 		canCreateDB: make(chan struct{}, o.maxDBs),
@@ -195,6 +221,10 @@ func New(ctx context.Context, opts ...Option) (*Env, error) {
 	}
 
 	err = waitForPostgresToBeReady(ctx, func(ctx context.Context) error {
+		// Don't try to connect until initialization is done. Postgres starts up twice
+		// in Docker, the first is just an initialization.
+		<-initDone
+
 		pool, err := pgxpool.Connect(ctx, env.dsn("" /* dbName */))
 		if err != nil {
 			return fmt.Errorf("failed to connect to database instance: %w", err)