@@ -32,6 +32,7 @@ import (
3232 "contrib.go.opencensus.io/exporter/prometheus"
3333 "contrib.go.opencensus.io/exporter/stackdriver"
3434 "github.com/GoogleCloudPlatform/alloydb-auth-proxy/alloydb"
35+ "github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/log"
3536 "github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/proxy"
3637 "github.com/spf13/cobra"
3738 "go.opencensus.io/trace"
@@ -76,8 +77,11 @@ func Execute() {
7677// Command represents an invocation of the AlloyDB Auth Proxy.
7778type Command struct {
7879 * cobra.Command
79- conf * proxy.Config
80+ conf * proxy.Config
81+ logger alloydb.Logger
82+ dialer alloydb.Dialer
8083
84+ cleanup func () error
8185 disableTraces bool
8286 telemetryTracingSampleRate int
8387 disableMetrics bool
@@ -88,25 +92,25 @@ type Command struct {
8892}
8993
9094// Option is a function that configures a Command.
91- type Option func (* proxy.Config )
95+ type Option func (* Command )
96+
97+ // WithLogger overrides the default logger.
98+ func WithLogger (l alloydb.Logger ) Option {
99+ return func (c * Command ) {
100+ c .logger = l
101+ }
102+ }
92103
93104// WithDialer configures the Command to use the provided dialer to connect to
94105// AlloyDB instances.
95106func WithDialer (d alloydb.Dialer ) Option {
96- return func (c * proxy. Config ) {
97- c .Dialer = d
107+ return func (c * Command ) {
108+ c .dialer = d
98109 }
99110}
100111
101112// NewCommand returns a Command object representing an invocation of the proxy.
102113func NewCommand (opts ... Option ) * Command {
103- c := & Command {
104- conf : & proxy.Config {},
105- }
106- for _ , o := range opts {
107- o (c .conf )
108- }
109-
110114 cmd := & cobra.Command {
111115 Use : "alloydb-auth-proxy instance_uri..." ,
112116 Version : versionString ,
@@ -115,19 +119,38 @@ func NewCommand(opts ...Option) *Command {
115119connecting to AlloyDB instances. It listens on a local port and forwards
116120connections to your instance's IP address, providing a secure connection
117121without having to manage any client SSL certificates.` ,
118- Args : func (cmd * cobra.Command , args []string ) error {
119- err := parseConfig (cmd , c .conf , args )
120- if err != nil {
121- return err
122- }
123- // The arguments are parsed. Usage is no longer needed.
124- cmd .SilenceUsage = true
125- return nil
126- },
127- RunE : func (* cobra.Command , []string ) error {
128- return runSignalWrapper (c )
122+ }
123+
124+ logger := log .NewStdLogger (os .Stdout , os .Stderr )
125+ c := & Command {
126+ Command : cmd ,
127+ logger : logger ,
128+ cleanup : func () error { return nil },
129+ conf : & proxy.Config {
130+ UserAgent : userAgent ,
129131 },
130132 }
133+ for _ , o := range opts {
134+ o (c )
135+ }
136+
137+ cmd .Args = func (cmd * cobra.Command , args []string ) error {
138+ // Handle logger separately from config
139+ if c .conf .StructuredLogs {
140+ c .logger , c .cleanup = log .NewStructuredLogger ()
141+ }
142+ err := parseConfig (c , c .conf , args )
143+ if err != nil {
144+ return err
145+ }
146+ // The arguments are parsed. Usage is no longer needed.
147+ cmd .SilenceUsage = true
148+ // Errors will be handled by logging from here on.
149+ cmd .SilenceErrors = true
150+ return nil
151+ }
152+
153+ cmd .RunE = func (* cobra.Command , []string ) error { return runSignalWrapper (c ) }
131154
132155 // Global-only flags
133156 cmd .PersistentFlags ().StringVarP (& c .conf .Token , "token" , "t" , "" ,
@@ -136,6 +159,8 @@ without having to manage any client SSL certificates.`,
136159 "Path to a service account key to use for authentication." )
137160 cmd .PersistentFlags ().BoolVarP (& c .conf .GcloudAuth , "gcloud-auth" , "g" , false ,
138161 "Use gcloud's user configuration to retrieve a token for authentication." )
162+ cmd .PersistentFlags ().BoolVarP (& c .conf .StructuredLogs , "structured-logs" , "l" , false ,
163+ "Enable structured logs using the LogEntry format" )
139164 cmd .PersistentFlags ().Uint64Var (& c .conf .MaxConnections , "max-connections" , 0 ,
140165 `Limits the number of connections by refusing any additional connections.
141166When this flag is not set, there is no limit.` )
@@ -168,18 +193,15 @@ the maximum time has passed. Defaults to 0s.`)
168193 cmd .PersistentFlags ().StringVarP (& c .conf .UnixSocket , "unix-socket" , "u" , "" ,
169194 `Enables Unix sockets for all listeners using the provided directory.` )
170195
171- c .Command = cmd
172196 return c
173197}
174198
175- func parseConfig (cmd * cobra. Command , conf * proxy.Config , args []string ) error {
199+ func parseConfig (cmd * Command , conf * proxy.Config , args []string ) error {
176200 // If no instance connection names were provided, error.
177201 if len (args ) == 0 {
178202 return newBadCommandError ("missing instance uri (e.g., /projects/$PROJECTS/locations/$LOCTION/clusters/$CLUSTER/instances/$INSTANCES)" )
179203 }
180204
181- conf .UserAgent = userAgent
182-
183205 userHasSet := func (f string ) bool {
184206 return cmd .PersistentFlags ().Lookup (f ).Changed
185207 }
@@ -210,13 +232,13 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
210232 }
211233
212234 if ! userHasSet ("telemetry-project" ) && userHasSet ("telemetry-prefix" ) {
213- cmd .Println ("Ignoring telementry-prefix as telemetry-project was not set" )
235+ cmd .logger . Infof ("Ignoring telementry-prefix as telemetry-project was not set" )
214236 }
215237 if ! userHasSet ("telemetry-project" ) && userHasSet ("disable-metrics" ) {
216- cmd .Println ("Ignoring disable-metrics as telemetry-project was not set" )
238+ cmd .logger . Infof ("Ignoring disable-metrics as telemetry-project was not set" )
217239 }
218240 if ! userHasSet ("telemetry-project" ) && userHasSet ("disable-traces" ) {
219- cmd .Println ("Ignoring disable-traces as telemetry-project was not set" )
241+ cmd .logger . Infof ("Ignoring disable-traces as telemetry-project was not set" )
220242 }
221243
222244 var ics []proxy.InstanceConnConfig
@@ -288,6 +310,7 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
288310
289311// runSignalWrapper watches for SIGTERM and SIGINT and interupts execution if necessary.
290312func runSignalWrapper (cmd * Command ) error {
313+ defer cmd .cleanup ()
291314 ctx , cancel := context .WithCancel (cmd .Context ())
292315 defer cancel ()
293316
@@ -340,7 +363,7 @@ func runSignalWrapper(cmd *Command) error {
340363 // Give the HTTP server a second to shutdown cleanly.
341364 ctx2 , _ := context .WithTimeout (context .Background (), time .Second )
342365 if err := server .Shutdown (ctx2 ); err != nil {
343- cmd .Printf ("failed to shutdown Prometheus HTTP server: %v\n " , err )
366+ cmd .logger . Errorf ("failed to shutdown Prometheus HTTP server: %v\n " , err )
344367 }
345368 }
346369 }()
@@ -378,7 +401,7 @@ func runSignalWrapper(cmd *Command) error {
378401 startCh := make (chan * proxy.Client )
379402 go func () {
380403 defer close (startCh )
381- p , err := proxy .NewClient (ctx , cmd .Command , cmd .conf )
404+ p , err := proxy .NewClient (ctx , cmd .dialer , cmd . logger , cmd .conf )
382405 if err != nil {
383406 shutdownCh <- fmt .Errorf ("unable to start: %v" , err )
384407 return
@@ -389,13 +412,15 @@ func runSignalWrapper(cmd *Command) error {
389412 var p * proxy.Client
390413 select {
391414 case err := <- shutdownCh :
415+ cmd .logger .Errorf ("The proxy has encountered a terminal error: %v" , err )
392416 return err
393417 case p = <- startCh :
394418 }
395- cmd .Println ("The proxy has started successfully and is ready for new connections!" )
419+ cmd .logger .Infof ("The proxy has started successfully and is ready for new connections!" )
420+ defer p .Close ()
396421 defer func () {
397422 if cErr := p .Close (); cErr != nil {
398- cmd .PrintErrf ( "The proxy failed to close cleanly : %v\n "cErr )
423+ cmd .logger . Errorf ( "error during shutdown : %v"cErr )
399424 }
400425 }()
401426
@@ -406,11 +431,11 @@ func runSignalWrapper(cmd *Command) error {
406431 err := <- shutdownCh
407432 switch {
408433 case errors .Is (err , errSigInt ):
409- cmd .PrintErrln ("SIGINT signal received. Shutting down..." )
434+ cmd .logger . Errorf ("SIGINT signal received. Shutting down..." )
410435 case errors .Is (err , errSigTerm ):
411- cmd .PrintErrln ("SIGTERM signal received. Shutting down..." )
436+ cmd .logger . Errorf ("SIGTERM signal received. Shutting down..." )
412437 default :
413- cmd .PrintErrf ("The proxy has encountered a terminal error: %v\n " , err )
438+ cmd .logger . Errorf ("The proxy has encountered a terminal error: %v" , err )
414439 }
415440 return err
416441}
0 commit comments