From bda74c49e67002090bb1a534eb0868e237f7bea2 Mon Sep 17 00:00:00 2001
From: Iheanyi Ekechukwu <iheanyi@planetscale.com>
Date: Thu, 22 Apr 2021 18:28:24 -0400
Subject: [PATCH 1/2] Style branches differently.

---
 internal/cmd/shell/shell.go | 7 +++++++
 internal/printer/printer.go | 5 +++++
 2 files changed, 12 insertions(+)

diff --git a/internal/cmd/shell/shell.go b/internal/cmd/shell/shell.go
index faddf707..66ee0e66 100644
--- a/internal/cmd/shell/shell.go
+++ b/internal/cmd/shell/shell.go
@@ -173,6 +173,13 @@ second argument:
 				"-P", port,
 			}
 
+			branchStyled := printer.BoldBlue(branch)
+			if branch == "main" {
+				branchStyled = printer.BoldRed(branch)
+			}
+			os.Setenv("MYSQL_PS1", printer.Bold(fmt.Sprintf("%s/%s> ", database, branchStyled)))
+			defer os.Unsetenv("MYSQL_PS1")
+
 			m := &mysql{}
 			err = m.Run(ctx, mysqlArgs...)
 			return err
diff --git a/internal/printer/printer.go b/internal/printer/printer.go
index 28794ab2..0bd9a554 100644
--- a/internal/printer/printer.go
+++ b/internal/printer/printer.go
@@ -229,6 +229,11 @@ func BoldBlue(msg string) string {
 	return color.New(color.FgBlue).Add(color.Bold).Sprint(msg)
 }
 
+// BoldRed returns a string formatted with red and bold.
+func BoldRed(msg string) string {
+	return color.New(color.FgRed).Add(color.Bold).Sprint(msg)
+}
+
 // Bold returns a string formatted with bold.
 func Bold(msg string) string {
 	// the 'color' package already handles IsTTY gracefully

From 0d0c80cf0c4bffbd8c6cfd2e7f2c53da61ebc9e2 Mon Sep 17 00:00:00 2001
From: Iheanyi Ekechukwu <iheanyi@planetscale.com>
Date: Thu, 22 Apr 2021 18:59:02 -0400
Subject: [PATCH 2/2] Use cmd.Env instead of native os.Setenv.

---
 internal/cmd/shell/shell.go | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/internal/cmd/shell/shell.go b/internal/cmd/shell/shell.go
index 66ee0e66..3f7e9e24 100644
--- a/internal/cmd/shell/shell.go
+++ b/internal/cmd/shell/shell.go
@@ -173,15 +173,9 @@ second argument:
 				"-P", port,
 			}
 
-			branchStyled := printer.BoldBlue(branch)
-			if branch == "main" {
-				branchStyled = printer.BoldRed(branch)
-			}
-			os.Setenv("MYSQL_PS1", printer.Bold(fmt.Sprintf("%s/%s> ", database, branchStyled)))
-			defer os.Unsetenv("MYSQL_PS1")
-
+			styledBranch := formatMySQLBranch(database, branch)
 			m := &mysql{}
-			err = m.Run(ctx, mysqlArgs...)
+			err = m.Run(ctx, styledBranch, mysqlArgs...)
 			return err
 
 		},
@@ -198,6 +192,16 @@ second argument:
 	return cmd
 }
 
+func formatMySQLBranch(database, branch string) string {
+	branchStyled := printer.BoldBlue(branch)
+	if branch == "main" {
+		branchStyled = printer.BoldRed(branch)
+	}
+
+	return printer.Bold(fmt.Sprintf("%s/%s> ", database, branchStyled))
+
+}
+
 // createLoginFile creates a temporary file to store the username and password, so we don't have to
 // pass them as `mysql` command-line arguments.
 func createLoginFile(username, password string) (string, error) {
@@ -219,12 +223,15 @@ type mysql struct {
 }
 
 // Run runs the `mysql` client with the given arguments.
-func (m *mysql) Run(ctx context.Context, args ...string) error {
+func (m *mysql) Run(ctx context.Context, styledBranch string, args ...string) error {
 	c := exec.CommandContext(ctx, "mysql", args...)
 	if m.Dir != "" {
 		c.Dir = m.Dir
 	}
 
+	c.Env = append(os.Environ(),
+		fmt.Sprintf("MYSQL_PS1=%s", styledBranch))
+
 	c.Stdout = os.Stdout
 	c.Stderr = os.Stderr
 	c.Stdin = os.Stdin