diff --git a/lib/client/db/postgres/repl/repl.go b/lib/client/db/postgres/repl/repl.go index 514d9160e3efb..66c8cb0f046f2 100644 --- a/lib/client/db/postgres/repl/repl.go +++ b/lib/client/db/postgres/repl/repl.go @@ -82,7 +82,7 @@ func New(_ context.Context, cfg *dbrepl.NewREPLConfig) (dbrepl.REPLInstance, err func (r *REPL) Run(ctx context.Context) error { pgConn, err := pgconn.ConnectConfig(ctx, r.connConfig) if err != nil { - return trace.Wrap(err) + return trace.ConnectionProblem(err, "Unable to connect to database: %v", err) } defer func() { closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) diff --git a/lib/client/db/postgres/repl/repl_test.go b/lib/client/db/postgres/repl/repl_test.go index 0aa03b84c8023..e2d780f9ab9c7 100644 --- a/lib/client/db/postgres/repl/repl_test.go +++ b/lib/client/db/postgres/repl/repl_test.go @@ -175,6 +175,19 @@ func TestClose(t *testing.T) { } } +func TestConnectionError(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + instance, tc := StartWithServer(t, ctx, WithSkipREPLRun()) + + // Force the server to be closed + tc.CloseServer() + + err := instance.Run(ctx) + require.Error(t, err) + require.True(t, trace.IsConnectionProblem(err), "expected run to be a connection error but got %T", err) +} + func writeLine(t *testing.T, c *testCtx, line string) { t.Helper() data := []byte(line + lineBreak) diff --git a/lib/web/databases.go b/lib/web/databases.go index f84beeb4ef7a8..57ad42e67b0d3 100644 --- a/lib/web/databases.go +++ b/lib/web/databases.go @@ -466,7 +466,6 @@ func (h *Handler) dbConnect( } stream := terminal.NewStream(ctx, terminal.StreamConfig{WS: ws}) - defer stream.Close() replConn, alpnConn := net.Pipe() sess := &databaseInteractiveSession{ @@ -486,11 +485,22 @@ func (h *Handler) dbConnect( } defer sess.Close() + // Don't close the terminal stream on session error, as it would also + // cause the underlying connection to be closed. This will prevent the + // middleware from properly writing the error into the WebSocket connection. if err := sess.Run(); err != nil { log.ErrorContext(ctx, "Database interactive session exited with error", "error", err) return nil, trace.Wrap(err) } + // TODO(gabrielcorado): Right now, if we send a close message the UI closes + // the terminal without giving the chance for users to review the session. + // Once this gets solved, we should send the close message here. + + if err := stream.Close(); err != nil { + log.ErrorContext(ctx, "Unable to close web socket terminal stream", "error", err) + } + return nil, nil } @@ -617,7 +627,7 @@ func (s *databaseInteractiveSession) Run() error { func (s *databaseInteractiveSession) Close() error { s.replConn.Close() - return s.ws.Close() + return nil } // issueCerts performs the MFA (if required) and generate the user session