Reply database REPL errors to the WebSocket client

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)

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)

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