feat: add support for wait command

To help ensure the Proxy is up and ready, this commit adds a wait
command with an optional `--max` flag to set the maximum time to wait.

By default when invoking this command:

./cloud-sql-proxy wait

The Proxy will wait up to the maximum time for the /startup endpoint to
respond. This command requires that the Proxy be started in another
process with the HTTP health check enabled. If an alternate health
check port or address is used, as in:

./cloud-sql-proxy <INSTANCE_CONNECTION_NAME> --http-address 0.0.0.0 \
  --http-port 9191

Then the wait command must also be told to use the same custom values:

./cloud-sql-proxy wait --http-address 0.0.0.0 \
  --http-port 9191

By default the wait command will wait 30 seconds. To alter this value,
use:

./cloud-sql-proxy wait --max 10s

Fixes #1117

cmd/root.go

@@ -280,8 +280,38 @@ Localhost Admin Server
     /quitquitquit. The admin server exits gracefully when it receives a POST
     request at /quitquitquit.
 
-(*) indicates a flag that may be used as a query parameter
+Waiting for Startup
+
+    Sometimes it is necessary to wait for the Proxy to start.
+
+    To help ensure the Proxy is up and ready, the Proxy includes adds a wait
+    subcommand with an optional --max flag to set the maximum time to wait.
+
+    By default when invoking this command:
+
+    ./cloud-sql-proxy wait
+
+    The Proxy will wait up to the maximum time for the /startup endpoint to
+    respond. This command requires that the Proxy be started in another
+    process with the HTTP health check enabled. If an alternate health
+    check port or address is used, as in:
+
+    ./cloud-sql-proxy <INSTANCE_CONNECTION_NAME> \
+      --http-address 0.0.0.0 \
+      --http-port 9191
+
+    Then the wait command must also be told to use the same custom values:
 
+    ./cloud-sql-proxy wait \
+      --http-address 0.0.0.0 \
+      --http-port 9191
+
+    By default the wait command will wait 30 seconds. To alter this value,
+    use:
+
+    ./cloud-sql-proxy wait --max 10s
+
+(*) indicates a flag that may be used as a query parameter
 `
 
 const envPrefix = "CSQL_PROXY"
@@ -314,9 +344,44 @@ func instanceFromEnv(args []string) []string {
 	return args
 }
 
+const (
+	waitMaxFlag     = "max"
+	httpAddressFlag = "http-address"
+	httpPortFlag    = "http-port"
+)
+
+func runWaitCmd(c *cobra.Command, args []string) error {
+	a, _ := c.Flags().GetString(httpAddressFlag)
+	p, _ := c.Flags().GetString(httpPortFlag)
+	addr := fmt.Sprintf("http://%v:%v/startup", a, p)
+
+	wait, err := c.Flags().GetDuration(waitMaxFlag)
+	if err != nil {
+		// This error should always be nil. If the error occurs, it means the
+		// wait flag name has changed where it was registered.
+		return err
+	}
+	c.SilenceUsage = true
+
+	t := time.After(wait)
+	for {
+		select {
+		case <-t:
+			return errors.New("command failed to complete successfully")
+		default:
+			resp, err := http.Get(addr)
+			if err != nil || resp.StatusCode != http.StatusOK {
+				time.Sleep(time.Second)
+				break
+			}
+			return nil
+		}
+	}
+}
+
 // NewCommand returns a Command object representing an invocation of the proxy.
 func NewCommand(opts ...Option) *Command {
-	cmd := &cobra.Command{
+	rootCmd := &cobra.Command{
 		Use:     "cloud-sql-proxy INSTANCE_CONNECTION_NAME...",
 		Version: versionString,
 		Short:   "cloud-sql-proxy authorizes and encrypts connections to Cloud SQL.",
@@ -325,7 +390,7 @@ func NewCommand(opts ...Option) *Command {
 
 	logger := log.NewStdLogger(os.Stdout, os.Stderr)
 	c := &Command{
-		Command: cmd,
+		Command: rootCmd,
 		logger:  logger,
 		cleanup: func() error { return nil },
 		conf: &proxy.Config{
@@ -336,7 +401,19 @@ func NewCommand(opts ...Option) *Command {
 		o(c)
 	}
 
-	cmd.Args = func(cmd *cobra.Command, args []string) error {
+	var waitCmd = &cobra.Command{
+		Use:  "wait",
+		RunE: runWaitCmd,
+	}
+	waitFlags := waitCmd.PersistentFlags()
+	waitFlags.DurationP(
+		waitMaxFlag, "m",
+		30*time.Second,
+		"maximum amount of time to wait for startup",
+	)
+	rootCmd.AddCommand(waitCmd)
+
+	rootCmd.Args = func(cmd *cobra.Command, args []string) error {
 		// If args is not already populated, try to read from the environment.
 		if len(args) == 0 {
 			args = instanceFromEnv(args)
@@ -358,9 +435,9 @@ func NewCommand(opts ...Option) *Command {
 		return nil
 	}
 
-	cmd.RunE = func(*cobra.Command, []string) error { return runSignalWrapper(c) }
+	rootCmd.RunE = func(*cobra.Command, []string) error { return runSignalWrapper(c) }
 
-	pflags := cmd.PersistentFlags()
+	pflags := rootCmd.PersistentFlags()
 
 	// Override Cobra's default messages.
 	pflags.BoolP("help", "h", false, "Display help information for cloud-sql-proxy")
@@ -405,9 +482,9 @@ the Proxy will then pick-up automatically.`)
 		"Enable Prometheus HTTP endpoint /metrics on localhost")
 	pflags.StringVar(&c.conf.PrometheusNamespace, "prometheus-namespace", "",
 		"Use the provided Prometheus namespace for metrics")
-	pflags.StringVar(&c.conf.HTTPAddress, "http-address", "localhost",
+	pflags.StringVar(&c.conf.HTTPAddress, httpAddressFlag, "localhost",
 		"Address for Prometheus and health check server")
-	pflags.StringVar(&c.conf.HTTPPort, "http-port", "9090",
+	pflags.StringVar(&c.conf.HTTPPort, httpPortFlag, "9090",
 		"Port for Prometheus and health check server")
 	pflags.BoolVar(&c.conf.Debug, "debug", false,
 		"Enable pprof on the localhost admin server")
@@ -432,7 +509,7 @@ https://cloud.google.com/storage/docs/requester-pays`)
 	pflags.StringVar(&c.conf.ImpersonationChain, "impersonate-service-account", "",
 		`Comma separated list of service accounts to impersonate. Last value
 is the target account.`)
-	cmd.PersistentFlags().BoolVar(&c.conf.Quiet, "quiet", false, "Log error messages only")
+	rootCmd.PersistentFlags().BoolVar(&c.conf.Quiet, "quiet", false, "Log error messages only")
 	pflags.BoolVar(&c.conf.AutoIP, "auto-ip", false,
 		`Supports legacy behavior of v1 and will try to connect to first IP
 address returned by the SQL Admin API. In most cases, this flag should not be used.

cmd/wait_test.go

@@ -0,0 +1,60 @@
+// Copyright 2023 Google LLC
+//
+// 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 cmd
+
+import (
+	"net"
+	"testing"
+)
+
+func TestWaitCommandFlags(t *testing.T) {
+	ln, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	host, port, err := net.SplitHostPort(ln.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))
+	}()
+
+	_, err = invokeProxyCommand([]string{
+		"wait",
+		"--http-address", host,
+		"--http-port", port,
+		"--max=1s",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestWaitCommandFails(t *testing.T) {
+	_, err := invokeProxyCommand([]string{
+		"wait",
+		// assuming default host and port
+		"--max=500ms",
+	})
+	if err == nil {
+		t.Fatal("wait should fail when endpoint does not respond")
+	}
+}