feat: add support for wait command

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 a wait
+    subcommand with an optional --max flag to set the maximum time to wait.
+
+    Invoke the wait command, like this:
+
+    ./cloud-sql-proxy wait
+
+    By default, the Proxy will wait up to the maximum time for the startup
+    endpoint to respond. The wait 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, _ []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,68 @@
+// 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 (
+	"io"
+	"net"
+	"testing"
+	"time"
+)
+
+func TestWaitCommandFlags(t *testing.T) {
+	ln, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	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()
+		// Use a read deadline to produce read error
+		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
+		// Read client request first.
+		io.ReadAll(conn)
+		// Write a generic 200 response back.
+		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=100ms",
+	})
+	if err == nil {
+		t.Fatal("wait should fail when endpoint does not respond")
+	}
+}