diff --git a/go.mod b/go.mod index b9ed822799..296af1ea24 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 9be106897b..02a4303e77 100644 --- a/go.sum +++ b/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= diff --git a/integration-tests/go-sql-server-driver/driver/cmd.go b/integration-tests/go-sql-server-driver/driver/cmd.go index 8e4756efe9..bc3703efe3 100644 --- a/integration-tests/go-sql-server-driver/driver/cmd.go +++ b/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 { diff --git a/integration-tests/go-sql-server-driver/driver/cmd_unix.go b/integration-tests/go-sql-server-driver/driver/cmd_unix.go index 6e795c1258..ba1d6be102 100644 --- a/integration-tests/go-sql-server-driver/driver/cmd_unix.go +++ b/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: diff --git a/integration-tests/go-sql-server-driver/driver/cmd_windows.go b/integration-tests/go-sql-server-driver/driver/cmd_windows.go index 4f8f0cb4e2..3a5cb46b13 100644 --- a/integration-tests/go-sql-server-driver/driver/cmd_windows.go +++ b/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 { diff --git a/server/server.go b/server/server.go index 25a03380da..8aadc7b33e 100644 --- a/server/server.go +++ b/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 }