From d8a38c51a01ac85a588ad003ecfb9de7cdcacb1f Mon Sep 17 00:00:00 2001
From: reginaldalfret <reginaldalfret@gmail.com>
Date: Mon, 26 Oct 2026 16:15:58 +0530
Subject: [PATCH] fix: preserve custom logger option and add
 WithStructuredLogging

Signed-off-by: reginaldalfret <reginaldalfret@gmail.com>
---
 cmd/options.go      |  8 ++++++++
 cmd/options_test.go | 33 +++++++++++++++++++++++++++++++++
 cmd/root.go         | 13 ++++++++-----
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/cmd/options.go b/cmd/options.go
index aceb01843..26828ab43 100644
--- a/cmd/options.go
+++ b/cmd/options.go
@@ -22,6 +22,7 @@ type Option func(*Command)
 // WithLogger overrides the default logger.
 func WithLogger(l cloudsql.Logger) Option {
 	return func(c *Command) {
+		c.customLogger = true
 		c.logger = l
 	}
 }
@@ -83,6 +84,13 @@ func WithQuietLogging() Option {
 	}
 }
 
+// WithStructuredLogging configures the Proxy to log messages using JSON.
+func WithStructuredLogging() Option {
+	return func(c *Command) {
+		c.conf.StructuredLogs = true
+	}
+}
+
 // WithDebugLogging configures the Proxy to log debug level messages.
 func WithDebugLogging() Option {
 	return func(c *Command) {
diff --git a/cmd/options_test.go b/cmd/options_test.go
index c3dc718f6..906cc52eb 100644
--- a/cmd/options_test.go
+++ b/cmd/options_test.go
@@ -134,6 +134,16 @@ func TestCommandOptions(t *testing.T) {
 			},
 			option: WithQuietLogging(),
 		},
+		{
+			desc: "with structured logging",
+			isValid: func(c *Command) error {
+				if !c.conf.StructuredLogs {
+					return errors.New("structured logs was false, but should be true")
+				}
+				return nil
+			},
+			option: WithStructuredLogging(),
+		},
 		{
 			desc: "with lazy refresh",
 			isValid: func(c *Command) error {
@@ -165,6 +175,7 @@ func TestCommandOptions(t *testing.T) {
 }
 
 func TestCommandOptionsOverridesCLI(t *testing.T) {
+	logger := log.NewStdLogger(io.Discard, io.Discard)
 	tcs := []struct {
 		desc    string
 		isValid func(*Command) error
@@ -193,6 +204,28 @@ func TestCommandOptionsOverridesCLI(t *testing.T) {
 			option: WithQuietLogging(),
 			args:   []string{"--quiet", "false"},
 		},
+		{
+			desc: "with logger overriding structured logs flag",
+			isValid: func(c *Command) error {
+				if c.logger != logger {
+					return errors.New("custom logger was overridden by structured logs flag")
+				}
+				return nil
+			},
+			option: WithLogger(logger),
+			args:   []string{"--structured-logs"},
+		},
+		{
+			desc: "with logger overriding quiet flag",
+			isValid: func(c *Command) error {
+				if c.logger != logger {
+					return errors.New("custom logger was overridden by quiet flag")
+				}
+				return nil
+			},
+			option: WithLogger(logger),
+			args:   []string{"--quiet"},
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
diff --git a/cmd/root.go b/cmd/root.go
index a5fd07f4c..c992dc0d0 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -91,6 +91,7 @@ type Command struct {
 	*cobra.Command
 	conf             *proxy.Config
 	logger           cloudsql.Logger
+	customLogger     bool
 	dialer           cloudsql.Dialer
 	cleanup          func() error
 	connRefuseNotify func()
@@ -663,11 +664,13 @@ func loadConfig(c *Command, args []string, opts []Option) error {
 		o(c)
 	}
 
-	// Handle logger separately from config
-	if c.conf.StructuredLogs {
-		c.logger = log.NewStructuredLogger(c.conf.Quiet)
-	} else if c.conf.Quiet {
-		c.logger = log.NewStdLogger(io.Discard, os.Stderr)
+	// Handle logger separately from config if not overridden by WithLogger
+	if !c.customLogger {
+		if c.conf.StructuredLogs {
+			c.logger = log.NewStructuredLogger(c.conf.Quiet)
+		} else if c.conf.Quiet {
+			c.logger = log.NewStdLogger(io.Discard, os.Stderr)
+		}
 	}
 
 	err = parseConfig(c, c.conf, args)