feat: add support for Cloud Monitoring and Cloud Trace

cmd/root.go

@@ -31,10 +31,12 @@ import (
 
 	"cloud.google.com/go/cloudsqlconn"
 	"contrib.go.opencensus.io/exporter/prometheus"
+	"contrib.go.opencensus.io/exporter/stackdriver"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/cloudsql"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/gcloud"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/proxy"
 	"github.com/spf13/cobra"
+	"go.opencensus.io/trace"
 	"golang.org/x/oauth2"
 )
 
@@ -67,8 +69,13 @@ type Command struct {
 	*cobra.Command
 	conf *proxy.Config
 
-	prometheusNamespace string
-	httpPort            string
+	telemetryTracing           bool
+	telemetryTracingSampleRate int
+	telemetryMetrics           bool
+	telemetryProject           string
+	telemetryPrefix            string
+	prometheusNamespace        string
+	httpPort                   string
 }
 
 // Option is a function that configures a Command.
@@ -120,6 +127,16 @@ any client SSL certificates.`,
 		"Path to a service account key to use for authentication.")
 	cmd.PersistentFlags().BoolVarP(&c.conf.GcloudAuth, "gcloud-auth", "g", false,
 		"Use gcloud's user configuration to retrieve a token for authentication.")
+	cmd.PersistentFlags().StringVar(&c.telemetryProject, "telemetry-project", "",
+		"Use the provided project ID for Cloud Monitoring and/or Cloud Trace integration.")
+	cmd.PersistentFlags().BoolVar(&c.telemetryTracing, "telemetry-traces", false,
+		"Enable Cloud Trace integration")
+	cmd.PersistentFlags().IntVar(&c.telemetryTracingSampleRate, "telemetry-sample-rate", 10_000,
+		"Configure the denominator of the probabilistic sample rate of traces sent to Cloud Trace\n(e.g., 10,000 traces 1/10,000 calls).")
+	cmd.PersistentFlags().BoolVar(&c.telemetryMetrics, "telemetry-metrics", false,
+		"Enable Cloud Monitoring integration")
+	cmd.PersistentFlags().StringVar(&c.telemetryPrefix, "telemetry-prefix", "",
+		"Prefix to use for Cloud Monitoring metrics.")
 	cmd.PersistentFlags().StringVar(&c.prometheusNamespace, "prometheus-namespace", "",
 		"Enable Prometheus for metric collection using the provided namespace")
 	cmd.PersistentFlags().StringVar(&c.httpPort, "http-port", "9090",
@@ -192,6 +209,19 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
 	}
 	conf.DialerOpts = opts
 
+	// If a user has set telemetry-project, but hasn't enabled one of metrics or
+	// traces, fail.
+	if userHasSet("telemetry-project") &&
+		!userHasSet("telemetry-metrics") &&
+		!userHasSet("telemetry-traces") {
+		return newBadCommandError("cannot specify --telemetry-project without also setting --telemetry-metrics or --telemetry-traces")
+	}
+	// If a user has enabled metrics or traces, but has not set the telemetry
+	// project, fail.
+	if (userHasSet("telemetry-metrics") || userHasSet("telemetry-traces")) &&
+		!userHasSet("telemetry-project") {
+		return newBadCommandError("must specify --telemetry-project when using --telemetry-metrics or --telemetry-traces")
+	}
 	if userHasSet("http-port") && !userHasSet("prometheus-namespace") {
 		return newBadCommandError("cannot specify --http-port without --prometheus-namespace")
 	}
@@ -268,6 +298,34 @@ func runSignalWrapper(cmd *Command) error {
 	ctx, cancel := context.WithCancel(cmd.Context())
 	defer cancel()
 
+	// Configure Cloud Trace and/or Cloud Monitoring based on command
+	// invocation. If a project has not been enabled, no traces or metrics are
+	// enabled.
+	if cmd.telemetryProject != "" && (cmd.telemetryTracing || cmd.telemetryMetrics) {
+		sd, err := stackdriver.NewExporter(stackdriver.Options{
+			ProjectID:    cmd.telemetryProject,
+			MetricPrefix: cmd.telemetryPrefix,
+		})
+		if err != nil {
+			return err
+		}
+		if cmd.telemetryMetrics {
+			err = sd.StartMetricsExporter()
+			if err != nil {
+				return err
+			}
+		}
+		if cmd.telemetryTracing {
+			s := trace.ProbabilitySampler(1 / float64(cmd.telemetryTracingSampleRate))
+			trace.ApplyConfig(trace.Config{DefaultSampler: s})
+			trace.RegisterExporter(sd)
+		}
+		defer func() {
+			sd.Flush()
+			sd.StopMetricsExporter()
+		}()
+	}
+
 	shutdownCh := make(chan error)
 
 	if cmd.prometheusNamespace != "" {

cmd/root_test.go

@@ -175,6 +175,18 @@ func TestNewCommandArguments(t *testing.T) {
 				}},
 			}),
 		},
+		{
+			desc: "enabling telemetry traces works",
+			args: []string{"--telemetry-project", "proj",
+				"--telemetry-traces", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{}),
+		},
+		{
+			desc: "enabling telemetry metrics works",
+			args: []string{"--telemetry-project", "proj",
+				"--telemetry-metrics", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -282,6 +294,21 @@ func TestNewCommandWithErrors(t *testing.T) {
 			args: []string{"proj:region:inst?unix-socket=/path&port=5000"},
 		},
 		{
+			desc: "when telemetry-project is set without metrics or tracing",
+			args: []string{"--telemetry-project", "proj", "proj:region:inst"},
+		},
+		{
+			desc: "when metrics are enabled without a telemetry project",
+			args: []string{"--telemetry-metrics", "proj:region:inst"},
+		},
+		{
+			desc: "when traces are enabled without a telemetry project",
+			args: []string{"--telemetry-traces", "proj:region:inst"},
+		},
+		{
+			desc: "when telemetry prefix is set without a telemetry project",
+			args: []string{"--telemetry-prefix", "myprefix", "--telemetry-traces", "proj:region:inst"},
+		},
 			desc: "enabling a Prometheus port without a namespace",
 			args: []string{"--htto-port", "1111", "proj:region:inst"},
 		},

go.mod

@@ -6,6 +6,7 @@ require (
 	cloud.google.com/go/cloudsqlconn v0.2.1-0.20220401153611-87e713b37755
 	cloud.google.com/go/compute v1.6.1
 	contrib.go.opencensus.io/exporter/prometheus v0.4.1
+	contrib.go.opencensus.io/exporter/stackdriver v0.13.13
 	github.com/GoogleCloudPlatform/cloudsql-proxy v1.29.0
 	github.com/coreos/go-systemd/v22 v22.3.2
 	github.com/denisenkom/go-mssqldb v0.12.0
@@ -15,6 +16,7 @@ require (
 	github.com/jackc/pgx/v4 v4.16.1
 	github.com/lib/pq v1.10.5
 	github.com/spf13/cobra v1.2.1
+	go.opencensus.io v0.23.0
 	go.uber.org/zap v1.21.0
 	golang.org/x/crypto v0.0.0-20220507011949-2cf3adece122 // indirect
 	golang.org/x/net v0.0.0-20220412020605-290c469a71a5