From b9403839a5a9901373e4651be006fa8c98314f47 Mon Sep 17 00:00:00 2001
From: "STeve (Xin) Huang" <xin.huang@goteleport.com>
Date: Wed, 5 Jul 2023 10:48:08 -0400
Subject: [PATCH] `tsh db connect` should prefer mongosh (#28668)

---
 lib/client/db/dbcmd/dbcmd.go      | 38 ++++++++++++++++++-------------
 lib/client/db/dbcmd/dbcmd_test.go | 22 +++++++++++++++---
 2 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/lib/client/db/dbcmd/dbcmd.go b/lib/client/db/dbcmd/dbcmd.go
index da7d46b339960..fb2563704e0de 100644
--- a/lib/client/db/dbcmd/dbcmd.go
+++ b/lib/client/db/dbcmd/dbcmd.go
@@ -375,40 +375,39 @@ func (c *CLICommandBuilder) getMySQLCommand() (*exec.Cmd, error) {
 	return c.getMySQLOracleCommand()
 }
 
+func (c *CLICommandBuilder) isBinAvailable(file string) bool {
+	_, err := c.options.exe.LookPath(file)
+	return err == nil
+}
+
 // isMariaDBBinAvailable returns true if "mariadb" binary is found in the system PATH.
 func (c *CLICommandBuilder) isMariaDBBinAvailable() bool {
-	_, err := c.options.exe.LookPath(mariadbBin)
-	return err == nil
+	return c.isBinAvailable(mariadbBin)
 }
 
 // isMySQLBinAvailable returns true if "mysql" binary is found in the system PATH.
 func (c *CLICommandBuilder) isMySQLBinAvailable() bool {
-	_, err := c.options.exe.LookPath(mysqlBin)
-	return err == nil
+	return c.isBinAvailable(mysqlBin)
 }
 
 // isMongoshBinAvailable returns true if "mongosh" binary is found in the system PATH.
 func (c *CLICommandBuilder) isMongoshBinAvailable() bool {
-	_, err := c.options.exe.LookPath(mongoshBin)
-	return err == nil
+	return c.isBinAvailable(mongoshBin)
 }
 
 // isElasticsearchSQLBinAvailable returns true if "elasticsearch-sql-cli" binary is found in the system PATH.
 func (c *CLICommandBuilder) isElasticsearchSQLBinAvailable() bool {
-	_, err := c.options.exe.LookPath(elasticsearchSQLBin)
-	return err == nil
+	return c.isBinAvailable(elasticsearchSQLBin)
 }
 
 // isOpenSearchCLIBinAvailable returns true if "opensearch-cli" binary is found in the system PATH.
 func (c *CLICommandBuilder) isOpenSearchCLIBinAvailable() bool {
-	_, err := c.options.exe.LookPath(openSearchCLIBin)
-	return err == nil
+	return c.isBinAvailable(openSearchCLIBin)
 }
 
 // isOpenSearchCLIBinAvailable returns true if "opensearchsql" binary is found in the system PATH.
 func (c *CLICommandBuilder) isOpenSearchSQLBinAvailable() bool {
-	_, err := c.options.exe.LookPath(openSearchSQLBin)
-	return err == nil
+	return c.isBinAvailable(openSearchSQLBin)
 }
 
 // isMySQLBinMariaDBFlavor checks if mysql binary comes from Oracle or MariaDB.
@@ -430,9 +429,16 @@ func (c *CLICommandBuilder) isMySQLBinMariaDBFlavor() (bool, error) {
 	return strings.Contains(strings.ToLower(string(mysqlVer)), "mariadb"), nil
 }
 
+func (c *CLICommandBuilder) shouldUseMongoshBin() bool {
+	// Use "mongosh" if available.
+	// If not, use legacy "mongo" if available.
+	// If both are not available, pick "mongosh" in print out.
+	return c.isMongoshBinAvailable() || !c.isBinAvailable(mongoBin)
+}
+
 func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 	// look for `mongosh`
-	hasMongosh := c.isMongoshBinAvailable()
+	useMongosh := c.shouldUseMongoshBin()
 
 	var args []string
 
@@ -447,7 +453,7 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 
 		var flags tlsFlags
 
-		if hasMongosh {
+		if useMongosh {
 			flags = tlsFlags{tls: "--tls", tlsCertKeyFile: "--tlsCertificateKeyFile", tlsCAFile: "--tlsCAFile"}
 		} else {
 			flags = tlsFlags{tls: "--ssl", tlsCertKeyFile: "--sslPEMKeyFile", tlsCAFile: "--sslCAFile"}
@@ -466,7 +472,7 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 			// mongosh does not load system CAs by default which will cause issues if
 			// the proxy presents a certificate signed by a non-recognized authority
 			// which your system trusts (e.g. mkcert).
-			if hasMongosh {
+			if useMongosh {
 				args = append(args, "--tlsUseSystemCA")
 			}
 		}
@@ -477,7 +483,7 @@ func (c *CLICommandBuilder) getMongoCommand() *exec.Cmd {
 	args = append(args, c.getMongoAddress())
 
 	// use `mongosh` if available
-	if hasMongosh {
+	if useMongosh {
 		return exec.Command(mongoshBin, args...)
 	}
 
diff --git a/lib/client/db/dbcmd/dbcmd_test.go b/lib/client/db/dbcmd/dbcmd_test.go
index ab872a7e63108..4c6eef1014ca6 100644
--- a/lib/client/db/dbcmd/dbcmd_test.go
+++ b/lib/client/db/dbcmd/dbcmd_test.go
@@ -317,7 +317,9 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 			dbProtocol:   defaults.ProtocolMongoDB,
 			databaseName: "mydb",
 			execer: &fakeExec{
-				execOutput: map[string][]byte{},
+				execOutput: map[string][]byte{
+					"mongo": []byte("legacy"),
+				},
 			},
 			cmd: []string{"mongo",
 				"--ssl",
@@ -327,12 +329,14 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 			wantErr: false,
 		},
 		{
-			name:         "mongodb no TLS",
+			name:         "mongodb no TLS (legacy)",
 			dbProtocol:   defaults.ProtocolMongoDB,
 			databaseName: "mydb",
 			opts:         []ConnectCommandFunc{WithNoTLS()},
 			execer: &fakeExec{
-				execOutput: map[string][]byte{},
+				execOutput: map[string][]byte{
+					"mongo": []byte("legacy"),
+				},
 			},
 			cmd: []string{"mongo",
 				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
@@ -387,6 +391,18 @@ func TestCLICommandBuilderGetConnectCommand(t *testing.T) {
 				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
 			},
 		},
+		{
+			name:         "mongosh preferred",
+			dbProtocol:   defaults.ProtocolMongoDB,
+			databaseName: "mydb",
+			opts:         []ConnectCommandFunc{WithNoTLS()},
+			execer: &fakeExec{
+				execOutput: map[string][]byte{}, // Cannot find either bin.
+			},
+			cmd: []string{"mongosh",
+				"mongodb://localhost:12345/mydb?serverSelectionTimeoutMS=5000",
+			},
+		},
 		{
 			name:         "sqlserver",
 			dbProtocol:   defaults.ProtocolSQLServer,