Fixed racy first-start database initialization

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/PuerkitoBio/goquery v1.8.1
 	github.com/cockroachdb/apd/v3 v3.2.3
 	github.com/cockroachdb/errors v1.7.5
-	github.com/dolthub/dolt/go v0.40.5-0.20260710160803-b101e404aa74
+	github.com/dolthub/dolt/go v0.40.5-0.20260710195005-67fd372af00e
 	github.com/dolthub/eventsapi_schema v0.0.0-20260310172945-37a9265ade69
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
 	github.com/dolthub/go-mysql-server v0.20.1-0.20260709224106-fecf6bb54eb0

go.sum

@@ -246,8 +246,8 @@ github.com/dolthub/aws-sdk-go-ini-parser v0.0.0-20250305001723-2821c37f6c12 h1:I
 github.com/dolthub/aws-sdk-go-ini-parser v0.0.0-20250305001723-2821c37f6c12/go.mod h1:rN7X8BHwkjPcfMQQ2QTAq/xM3leUSGLfb+1Js7Y6TVo=
 github.com/dolthub/dolt-mcp v0.3.4 h1:AyG5cw+fNWXDHXujtQnqUPZrpWtPg6FN6yYtjv1pP44=
 github.com/dolthub/dolt-mcp v0.3.4/go.mod h1:bCZ7KHvDYs+M0e+ySgmGiNvLhcwsN7bbf5YCyillLrk=
-github.com/dolthub/dolt/go v0.40.5-0.20260710160803-b101e404aa74 h1:7sMiWh9QW8mdvlO5VoZ2Kh/QABlOpXRmxumqwVscCiA=
-github.com/dolthub/dolt/go v0.40.5-0.20260710160803-b101e404aa74/go.mod h1:9BHm5IeGCRju3T86BtqT671Iv189t9A4JZwy/rpu7cA=
+github.com/dolthub/dolt/go v0.40.5-0.20260710195005-67fd372af00e h1:oqowLp7vIKiI8WRcXabfeYSijJs6ATs+xv+SZv/cSwY=
+github.com/dolthub/dolt/go v0.40.5-0.20260710195005-67fd372af00e/go.mod h1:9BHm5IeGCRju3T86BtqT671Iv189t9A4JZwy/rpu7cA=
 github.com/dolthub/eventsapi_schema v0.0.0-20260310172945-37a9265ade69 h1:JShhbqMw26nKx3pqqu/cFxOpzBkN+4elVhzuUfgDw2k=
 github.com/dolthub/eventsapi_schema v0.0.0-20260310172945-37a9265ade69/go.mod h1:SSLraQS/jGLYFgff3vuZ+JbVUct6vyEeMzjLBqWqoyM=
 github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2 h1:u3PMzfF8RkKd3lB9pZ2bfn0qEG+1Gms9599cr0REMww=

integration-tests/go-sql-server-driver/driver/cmd.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -206,26 +207,79 @@ func (rs RepoStore) initDatabase(name string, fn func(db *sql.DB) error) error {
 	if err := cmd.Start(); err != nil {
 		return err
 	}
+
+	stopped := false
+	stopServer := func() error {
+		if stopped {
+			return nil
+		}
+		stopped = true
+		if err := interruptCmd(cmd); err != nil && !errors.Is(err, os.ErrProcessDone) {
+			return fmt.Errorf("could not interrupt init server for %s: %w (output: %s)", name, err, output.String())
+		}
+		if err := cmd.Wait(); err != nil {
+			return fmt.Errorf("init server for %s did not exit cleanly: %w (output: %s)", name, err, output.String())
+		}
+		return nil
+	}
+	// Best-effort shutdown on error paths; the success path below calls
+	// stopServer explicitly and checks its result.
 	defer func() {
-		_ = cmd.Process.Signal(os.Interrupt)
-		_, _ = cmd.Process.Wait()
+		_ = stopServer()
 	}()
 
-	db, err := ConnectDB("postgres", "password", name, "127.0.0.1", port, nil)
+	db, err := ConnectDB("postgres", "password", "", "127.0.0.1", port, nil)
 	if err != nil {
 		return fmt.Errorf("could not connect to init server for %s: %w (output: %s)", name, err, output.String())
 	}
-	defer db.Close()
-
-	if _, err := db.Exec(fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, name)); err != nil {
+	_, err = db.Exec(fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, name))
+	db.Close()
+	if err != nil {
 		return err
 	}
+
+	// Connect to the new database, wait for its initialization to fully
+	// complete, and run any additional setup.
+	db, err = ConnectDB("postgres", "password", name, "127.0.0.1", port, nil)
+	if err != nil {
+		return fmt.Errorf("could not connect to database %s on init server: %w (output: %s)", name, err, output.String())
+	}
+	defer db.Close()
+	if err := waitForDatabaseInit(db); err != nil {
+		return fmt.Errorf("database %s was not fully initialized: %w (output: %s)", name, err, output.String())
+	}
+
+	// Complete any requested setup
 	if fn != nil {
 		if err := fn(db); err != nil {
 			return err
 		}
 	}
-	return nil
+
+	// Finally shutdown the server now that initialization is complete
+	return stopServer()
+}
+
+// waitForDatabaseInit blocks until the database that |db| is connected to has
+// been fully initialized.
+func waitForDatabaseInit(db *sql.DB) error {
+	var lastErr error
+	for i := 0; i < ConnectAttempts; i++ {
+		if i != 0 {
+			time.Sleep(RetrySleepDuration)
+		}
+		var n int
+		err := db.QueryRow(`SELECT count(*) FROM dolt_log WHERE message = 'CREATE DATABASE'`).Scan(&n)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		if n > 0 {
+			return nil
+		}
+		lastErr = errors.New("no CREATE DATABASE commit found in dolt_log")
+	}
+	return lastErr
 }
 
 func sanitize(s string) string {

integration-tests/go-sql-server-driver/driver/cmd_unix.go

@@ -18,6 +18,7 @@
 package driver
 
 import (
+	"os"
 	"os/exec"
 	"syscall"
 )
@@ -26,6 +27,12 @@ func ApplyCmdAttributes(cmd *exec.Cmd) {
 	// nothing to do on unix / darwin
 }
 
+// interruptCmd requests that |cmd| shut down gracefully, the same way a user
+// pressing Ctrl-C would.
+func interruptCmd(cmd *exec.Cmd) error {
+	return cmd.Process.Signal(os.Interrupt)
+}
+
 func (s *SqlServer) GracefulStop() error {
 	select {
 	case <-s.Done:

integration-tests/go-sql-server-driver/driver/cmd_windows.go

@@ -29,6 +29,14 @@ func ApplyCmdAttributes(cmd *exec.Cmd) {
 	}
 }
 
+// interruptCmd requests that |cmd| shut down gracefully, the same way a user
+// pressing Ctrl-C would. os.Interrupt cannot be sent with Process.Signal on
+// Windows, so a console ctrl-break event is used instead (the process was
+// placed in its own process group by ApplyCmdAttributes).
+func interruptCmd(cmd *exec.Cmd) error {
+	return windows.GenerateConsoleCtrlEvent(windows.CTRL_BREAK_EVENT, uint32(cmd.Process.Pid))
+}
+
 func (s *SqlServer) GracefulStop() error {
 	err := windows.GenerateConsoleCtrlEvent(windows.CTRL_BREAK_EVENT, uint32(s.Cmd.Process.Pid))
 	if err != nil {

server/server.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
+	"github.com/dolthub/dolt/go/cmd/dolt/commands/engine"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/sqlserver"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
@@ -37,7 +38,6 @@ import (
 	"github.com/dolthub/dolt/go/store/util/tempfiles"
 	"github.com/dolthub/go-mysql-server/server"
 	"github.com/dolthub/go-mysql-server/sql"
-	"github.com/jackc/pgx/v5"
 
 	"github.com/dolthub/doltgresql/server/auth"
 	"github.com/dolthub/doltgresql/server/initialization"
@@ -148,6 +148,15 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 		return true, nil
 	})
 
+	// When we need to create the default database, do so via an engine initializer, which runs after the engine is
+	// constructed but before the server accepts any connections. This is necessary due to the fact that
+	// `CREATE DATABASE` is non-transactional and non-atomic, so other clients could connect to a partially-created
+	// database otherwise. This problem exists in Dolt as well, but Dolt doesn't auto-create a database on init.
+	var engineInitializer sqlserver.EngineInitializer
+	if initializeDefaultDatabase {
+		engineInitializer = defaultDatabaseInitializer{}
+	}
+
 	controller := svcs.NewController()
 	newCtx, cancelF := context.WithCancel(ctx)
 	go func() {
@@ -171,6 +180,7 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 		DoltEnv:                 dEnv,
 		ProtocolListenerFactory: protocolListenerFactory,
 		ProviderFactory:         DoltgresProviderFactory{},
+		EngineInitializer:       engineInitializer,
 	})
 	go controller.Start(newCtx)
 
@@ -179,39 +189,44 @@ func runServer(ctx context.Context, cfg *servercfg.DoltgresConfig, dEnv *env.Dol
 		return nil, err
 	}
 
-	if initializeDefaultDatabase {
-		err = createDefaultDatabase(ssCfg)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	// TODO: shutdown replication cleanly when we stop the server
 	_, err = startReplication(cfg, ssCfg)
 	if err != nil {
+		controller.Stop()
 		return nil, err
 	}
 
 	return controller, nil
 }
 
-// createDefaultDatabase creates the database named on the local server using the configuration values to connect, returning
-// any error
-func createDefaultDatabase(cfg doltservercfg.ServerConfig) error {
-	user, password := auth.GetSuperUserAndPassword()
-	dbName := getDefaultDatabaseName(user)
+// defaultDatabaseInitializer creates the default database on a first run against an empty data dir. It runs via
+// sqlserver.Config.EngineInitializer, after the engine has been constructed but before the server accepts any
+// connections, so no client can observe (or interfere with) a partially created default database.
+type defaultDatabaseInitializer struct{}
 
-	dsn := fmt.Sprintf("postgres://%s:%s@localhost:%d", user, password, cfg.Port())
+var _ sqlserver.EngineInitializer = defaultDatabaseInitializer{}
 
-	// Connect to the server and create the default database with the given name.
-	ctx := context.Background()
-	conn, err := pgx.Connect(ctx, dsn)
+// InitializeEngine implements sqlserver.EngineInitializer.
+func (defaultDatabaseInitializer) InitializeEngine(ctx context.Context, se *engine.SqlEngine) error {
+	user, _ := auth.GetSuperUserAndPassword()
+	dbName := getDefaultDatabaseName(user)
+
+	sqlCtx, err := se.NewDefaultContext(ctx)
 	if err != nil {
 		return err
 	}
-	defer conn.Close(ctx)
-
-	_, err = conn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s;", dbName))
+	// The session must run as the doltgres superuser: unlike Dolt, doltgres has
+	// no "root" user, so NewLocalContext's default client would be rejected.
+	sqlCtx.Session.SetClient(sql.Client{User: user, Address: "localhost", Capabilities: 0})
+	defer sql.SessionEnd(sqlCtx.Session)
+	sql.SessionCommandBegin(sqlCtx.Session)
+	defer sql.SessionCommandEnd(sqlCtx.Session)
+
+	_, iter, _, err := se.Query(sqlCtx, fmt.Sprintf("CREATE DATABASE %s;", dbName))
+	if err != nil {
+		return err
+	}
+	_, err = sql.RowIterToRows(sqlCtx, iter)
 	return err
 }