[v10] Mongo clients with serverSelectionTimeoutMS set to 5000

lib/client/db/dbcmd/dbcmd.go

@@ -20,22 +20,25 @@ package dbcmd
 
 import (
 	"fmt"
+	"net/url"
+	"os"
 	"os/exec"
 	"path/filepath"
 	"strconv"
 	"strings"
 
+	"github.com/gravitational/trace"
+	"github.com/sirupsen/logrus"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
+
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/client/db"
 	"github.com/gravitational/teleport/lib/client/db/mysql"
 	"github.com/gravitational/teleport/lib/client/db/postgres"
 	"github.com/gravitational/teleport/lib/defaults"
+	libdbcommon "github.com/gravitational/teleport/lib/srv/db/common"
 	"github.com/gravitational/teleport/lib/tlsca"
 	"github.com/gravitational/teleport/lib/utils"
-
-	"github.com/gravitational/trace"
-
-	"github.com/sirupsen/logrus"
 )
 
 const (
@@ -359,10 +362,7 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 	// look for `mongosh`
 	hasMongosh := c.isMongoshBinAvailable()
 
-	args := []string{
-		"--host", c.host,
-		"--port", strconv.Itoa(c.port),
-	}
+	var args []string
 
 	if !c.options.noTLS {
 		// Starting with Mongo 4.2 there is an updated set of flags.
@@ -400,9 +400,9 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 		}
 	}
 
-	if c.db.Database != "" {
-		args = append(args, c.db.Database)
-	}
+	// Add the address at the end. Address contains host, port, database name,
+	// and other options like server selection timeout.
+	args = append(args, c.getMongoAddress())
 
 	// use `mongosh` if available
 	if hasMongosh {
@@ -413,6 +413,33 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 	return c.exe.Command(mongoBin, args...)
 }
 
+func (c *CLICommandBuilder) getMongoAddress() string {
+	query := make(url.Values)
+
+	// Use the same default server selection timeout that the backend engine is
+	// using. The environment variable serves as a hidden option to force a
+	// different timeout for debugging purpose or extreme situations.
+	serverSelectionTimeoutMS := strconv.Itoa(int(libdbcommon.DefaultMongoDBServerSelectionTimeout.Milliseconds()))
+	if envValue := os.Getenv(envVarMongoServerSelectionTimeoutMS); envValue != "" {
+		c.options.log.Infof("Using environment variable %s=%s.", envVarMongoServerSelectionTimeoutMS, envValue)
+		serverSelectionTimeoutMS = envValue
+	}
+	query.Set("serverSelectionTimeoutMS", serverSelectionTimeoutMS)
+
+	address := url.URL{
+		Scheme:   connstring.SchemeMongoDB,
+		Host:     fmt.Sprintf("%s:%d", c.host, c.port),
+		RawQuery: query.Encode(),
+		Path:     fmt.Sprintf("/%s", c.db.Database),
+	}
+
+	// Quote the address for printing as the address contains "?".
+	if c.options.printFormat {
+		return fmt.Sprintf(`"%s"`, address.String())
+	}
+	return address.String()
+}
+
 // getRedisCommand returns redis-cli commands used by 'tsh db connect' when connecting to a Redis instance.
 func (c *CLICommandBuilder) getRedisCommand() *exec.Cmd {
 	// TODO(jakub): Add "-3" when Teleport adds support for Redis RESP3 protocol.
@@ -519,6 +546,17 @@ func WithNoTLS() ConnectCommandFunc {
 
 // WithPrintFormat is the connect command option that hints the command will be
 // printed instead of being executed.
+//
+// For example, when enabled, a quote will be used for Postgres and MongoDB
+// connection strings to avoid "&" getting interpreted by the shell.
+//
+// WithPrintFormat is known to be used for the following situations:
+// - tsh db config --format cmd <database>
+// - tsh proxy db --tunnel <database>
+// - Teleport Connect where the command is put into a terminal.
+//
+// WithPrintFormat should NOT be used when the exec.Cmd gets executed by the
+// client application.
 func WithPrintFormat() ConnectCommandFunc {
 	return func(opts *connectionCommandOpts) {
 		opts.printFormat = true
@@ -548,3 +586,9 @@ func WithTolerateMissingCLIClient() ConnectCommandFunc {
 		opts.tolerateMissingCLIClient = true
 	}
 }
+
+const (
+	// envVarMongoServerSelectionTimeoutMS is the environment variable that
+	// controls the server selection timeout used for MongoDB clients.
+	envVarMongoServerSelectionTimeoutMS = "TELEPORT_MONGO_SERVER_SELECTION_TIMEOUT_MS"
+)

lib/client/db/dbcmd/dbcmd_test.go

@@ -336,11 +336,10 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				execOutput: map[string][]byte{},
 			},
 			cmd: []string{"mongo",
-				"--host", "localhost",
-				"--port", "12345",
 				"--ssl",
 				"--sslPEMKeyFile", "/tmp/keys/example.com/bob-db/db.example.com/mysql-x509.pem",
-				"mydb"},
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
 			wantErr: false,
 		},
 		{
@@ -352,13 +351,12 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				execOutput: map[string][]byte{},
 			},
 			cmd: []string{"mongo",
-				"--host", "localhost",
-				"--port", "12345",
-				"mydb"},
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
 			wantErr: false,
 		},
 		{
-			name:         "mongosh",
+			name:         "mongosh no CA",
 			dbProtocol:   defaults.ProtocolMongoDB,
 			databaseName: "mydb",
 			execer: &fakeExec{
@@ -367,12 +365,11 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				},
 			},
 			cmd: []string{"mongosh",
-				"--host", "localhost",
-				"--port", "12345",
 				"--tls",
 				"--tlsCertificateKeyFile", "/tmp/keys/example.com/bob-db/db.example.com/mysql-x509.pem",
 				"--tlsUseSystemCA",
-				"mydb"},
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
 		},
 		{
 			name:         "mongosh",
@@ -386,12 +383,11 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				},
 			},
 			cmd: []string{"mongosh",
-				"--host", "localhost",
-				"--port", "12345",
 				"--tls",
 				"--tlsCertificateKeyFile", "/tmp/keys/example.com/bob-db/db.example.com/mysql-x509.pem",
 				"--tlsCAFile", "/tmp/keys/example.com/cas/example.com.pem",
-				"mydb"},
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
 		},
 		{
 			name:         "mongosh no TLS",
@@ -404,9 +400,8 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				},
 			},
 			cmd: []string{"mongosh",
-				"--host", "localhost",
-				"--port", "12345",
-				"mydb"},
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
 		},
 		{
 			name:         "sqlserver",

lib/srv/db/common/constants.go

@@ -0,0 +1,25 @@
+/*
+Copyright 2022 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package common
+
+import "time"
+
+const (
+	// DefaultMongoDBServerSelectionTimeout is the timeout for selecting a
+	// MongoDB server to connect to.
+	DefaultMongoDBServerSelectionTimeout = 5 * time.Second
+)

lib/srv/db/mongodb/connect.go

@@ -98,7 +98,7 @@ func (e *Engine) getTopologyOptions(ctx context.Context, sessionCtx *common.Sess
 			return connString
 		}),
 		topology.WithServerSelectionTimeout(func(time.Duration) time.Duration {
-			return serverSelectionTimeout
+			return common.DefaultMongoDBServerSelectionTimeout
 		}),
 		topology.WithServerOptions(func(so ...topology.ServerOption) []topology.ServerOption {
 			return serverOptions
@@ -202,9 +202,3 @@ func (h *handshaker) GetHandshakeInformation(context.Context, address.Address, d
 func (h *handshaker) FinishHandshake(context.Context, driver.Connection) error {
 	return nil
 }
-
-const (
-	// serverSelectionTimeout is the timeout for selecting a MongoDB server
-	// to connect to.
-	serverSelectionTimeout = 5 * time.Second
-)

lib/teleterm/clusters/cluster_gateways.go

@@ -97,6 +97,7 @@ func buildCLICommand(c *Cluster, gw *gateway.Gateway) (*exec.Cmd, error) {
 		dbcmd.WithLogger(gw.Log),
 		dbcmd.WithLocalProxy(gw.LocalAddress, gw.LocalPortInt(), ""),
 		dbcmd.WithNoTLS(),
+		dbcmd.WithPrintFormat(),
 		dbcmd.WithTolerateMissingCLIClient(),
 	).GetConnectCommandNoAbsPath()
 

tool/tsh/proxy.go

@@ -298,6 +298,7 @@ func onProxyCommandDB(cf *CLIConf) error {
 			dbcmd.WithLocalProxy("localhost", addr.Port(0), ""),
 			dbcmd.WithNoTLS(),
 			dbcmd.WithLogger(log),
+			dbcmd.WithPrintFormat(),
 		).GetConnectCommand()
 		if err != nil {
 			return trace.Wrap(err)