msp/runtime: export Cloud SQL conection internals for direct usage (#62524)

Part of https://linear.app/sourcegraph/issue/CORE-96 - we want to be able to use the same Cloud SQL connection mechanism we use for MSP-provisioned databases to be able to connect to a database in another project (in this case, the database of interest is the Sourcegraph.com database).

Also see #62525 

## Test plan

CI - this is a refactor only

lib/managedservicesplatform/cloudsql/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "cloudsql",
+    srcs = ["cloudsql.go"],
+    importpath = "github.com/sourcegraph/sourcegraph/lib/managedservicesplatform/cloudsql",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//lib/errors",
+        "@com_github_jackc_pgx_v5//:pgx",
+        "@com_github_jackc_pgx_v5//stdlib",
+        "@com_google_cloud_go_cloudsqlconn//:cloudsqlconn",
+    ],
+)

lib/managedservicesplatform/cloudsql/cloudsql.go

@@ -0,0 +1,93 @@
+package cloudsql
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"net"
+
+	"cloud.google.com/go/cloudsqlconn"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/stdlib"
+
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+)
+
+type ConnConfig struct {
+	// ConnectionName is the CloudSQL connection name,
+	// e.g. '${project}:${region}:${instance}'
+	ConnectionName *string
+	// User is the Cloud SQL user to connect as, e.g. 'test-sa@test-project.iam'
+	User *string
+	// Database to connect to.
+	Database string
+	// DialOptions are any additional options to pass to the underlying
+	// cloud-sql-proxy driver.
+	DialOptions []cloudsqlconn.DialOption
+}
+
+// Open opens a *sql.DB connection to the Cloud SQL instance specified by the
+// ConnConfig.
+//
+// 🔔 If you are connecting to a MSP-provisioned Cloud SQL instance,
+// DO NOT use this - instead, use runtime.Contract.PostgreSQL.OpenDatabase
+// instead.
+func Open(
+	ctx context.Context,
+	cfg ConnConfig,
+) (*sql.DB, error) {
+	config, err := getCloudSQLConnConfig(ctx, cfg)
+	if err != nil {
+		return nil, errors.Wrap(err, "get CloudSQL connection config")
+	}
+	return sql.Open("pgx", stdlib.RegisterConnConfig(config))
+}
+
+// Connect opens a *pgx.Conn connection to the CloudSQL instance specified by
+// the ConnConfig.
+//
+// 🔔 If you are connecting to a MSP-provisioned Cloud SQL instance,
+// DO NOT use this - instead, use runtime.Contract.PostgreSQL.OpenDatabase
+// instead.
+func Connect(
+	ctx context.Context,
+	cfg ConnConfig,
+) (*pgx.Conn, error) {
+	config, err := getCloudSQLConnConfig(ctx, cfg)
+	if err != nil {
+		return nil, errors.Wrap(err, "get CloudSQL connection config")
+	}
+	return pgx.ConnectConfig(ctx, config)
+}
+
+// getCloudSQLConnConfig generates a pgx connection configuration for using
+// a Cloud SQL instance using IAM auth.
+func getCloudSQLConnConfig(
+	ctx context.Context,
+	cfg ConnConfig,
+) (*pgx.ConnConfig, error) {
+	if cfg.ConnectionName == nil || cfg.User == nil {
+		return nil, errors.New("missing required PostgreSQL configuration")
+	}
+
+	// https://github.com/GoogleCloudPlatform/cloud-sql-go-connector?tab=readme-ov-file#automatic-iam-database-authentication
+	dsn := fmt.Sprintf("user=%s dbname=%s", *cfg.User, cfg.Database)
+	config, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		return nil, errors.Wrap(err, "pgx.ParseConfig")
+	}
+	customDialer, err := cloudsqlconn.NewDialer(ctx,
+		// always the case when using Cloud SQL in MSP
+		cloudsqlconn.WithIAMAuthN(),
+		// allow passthrough of additional dial options
+		cloudsqlconn.WithDefaultDialOptions(cfg.DialOptions...))
+	if err != nil {
+		return nil, errors.Wrap(err, "cloudsqlconn.NewDialer")
+	}
+	// Use the Cloud SQL connector to handle connecting to the instance.
+	// This approach does *NOT* require the Cloud SQL proxy.
+	config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
+		return customDialer.Dial(ctx, *cfg.ConnectionName)
+	}
+	return config, nil
+}

lib/managedservicesplatform/runtime/contract/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
     deps = [
         "//lib/errors",
         "//lib/managedservicesplatform/bigquerywriter",
+        "//lib/managedservicesplatform/cloudsql",
         "//lib/managedservicesplatform/runtime/internal/opentelemetry",
         "//lib/pointers",
         "@com_github_getsentry_sentry_go//:sentry-go",

lib/managedservicesplatform/runtime/contract/postgresql.go

@@ -4,14 +4,14 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
-	"fmt"
-	"net"
 	"text/template"
 
 	"cloud.google.com/go/cloudsqlconn"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/stdlib"
+
 	"github.com/sourcegraph/sourcegraph/lib/errors"
+	"github.com/sourcegraph/sourcegraph/lib/managedservicesplatform/cloudsql"
 )
 
 type postgreSQLContract struct {
@@ -45,48 +45,56 @@ func (c postgreSQLContract) Configured() bool {
 // variable.
 func (c postgreSQLContract) OpenDatabase(ctx context.Context, database string) (*sql.DB, error) {
 	if c.customDSNTemplate != nil {
-		tmpl, err := template.New("PGDSN").Parse(*c.customDSNTemplate)
+		config, err := parseCustomDSNTemplateConnConfig(*c.customDSNTemplate, database)
 		if err != nil {
-			return nil, errors.Wrap(err, "PGDSN is not a valid template")
+			return nil, err
 		}
-		var dsn bytes.Buffer
-		if err := tmpl.Execute(&dsn, struct{ Database string }{Database: database}); err != nil {
-			return nil, errors.Wrap(err, "PGDSN template is invalid")
-		}
-		return sql.Open("pgx", dsn.String())
+		return sql.Open("customdsn", stdlib.RegisterConnConfig(config))
 	}
-
-	config, err := c.getCloudSQLConnConfig(ctx, database)
-	if err != nil {
-		return nil, errors.Wrap(err, "get CloudSQL connection config")
-	}
-	return sql.Open("pgx", stdlib.RegisterConnConfig(config))
+	return cloudsql.Open(ctx, c.getCloudSQLConnConfig(database))
 }
 
-// getCloudSQLConnConfig generates a pgx connection configuration for using
-// a Cloud SQL instance using IAM auth.
-func (c postgreSQLContract) getCloudSQLConnConfig(ctx context.Context, database string) (*pgx.ConnConfig, error) {
-	if c.instanceConnectionName == nil || c.instanceConnectionUser == nil {
-		return nil, errors.New("missing required PostgreSQL configuration")
+// ConnectToDatabase is similar to OpenDatabase, but returns a
+// github.com/jackc/pgx/v5 connection to the configured datbase instead for
+// services that prefer to use 'pgx' directly.
+//
+// In development, the connection can be overridden with the PGDSN environment
+// variable.
+func (c postgreSQLContract) ConnectToDatabase(ctx context.Context, database string) (*pgx.Conn, error) {
+	if c.customDSNTemplate != nil {
+		config, err := parseCustomDSNTemplateConnConfig(*c.customDSNTemplate, database)
+		if err != nil {
+			return nil, err
+		}
+		return pgx.ConnectConfig(ctx, config)
 	}
+	return cloudsql.Connect(ctx, c.getCloudSQLConnConfig(database))
+}
 
-	// https://github.com/GoogleCloudPlatform/cloud-sql-go-connector?tab=readme-ov-file#automatic-iam-database-authentication
-	dsn := fmt.Sprintf("user=%s dbname=%s", *c.instanceConnectionUser, database)
-	config, err := pgx.ParseConfig(dsn)
-	if err != nil {
-		return nil, errors.Wrap(err, "pgx.ParseConfig")
+func (c postgreSQLContract) getCloudSQLConnConfig(database string) cloudsql.ConnConfig {
+	return cloudsql.ConnConfig{
+		ConnectionName: c.instanceConnectionName,
+		User:           c.instanceConnectionUser,
+		Database:       database,
+		DialOptions: []cloudsqlconn.DialOption{
+			// MSP-provisioned databases only allow private IP access
+			cloudsqlconn.WithPrivateIP(),
+		},
 	}
-	d, err := cloudsqlconn.NewDialer(ctx,
-		cloudsqlconn.WithIAMAuthN(),
-		// MSP uses private IP
-		cloudsqlconn.WithDefaultDialOptions(cloudsqlconn.WithPrivateIP()))
+}
+
+func parseCustomDSNTemplateConnConfig(customDSNTemplate, database string) (*pgx.ConnConfig, error) {
+	tmpl, err := template.New("PGDSN").Parse(customDSNTemplate)
 	if err != nil {
-		return nil, errors.Wrap(err, "cloudsqlconn.NewDialer")
+		return nil, errors.Wrap(err, "PGDSN is not a valid template")
 	}
-	// Use the Cloud SQL connector to handle connecting to the instance.
-	// This approach does *NOT* require the Cloud SQL proxy.
-	config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
-		return d.Dial(ctx, *c.instanceConnectionName)
+	var dsn bytes.Buffer
+	if err := tmpl.Execute(&dsn, struct{ Database string }{Database: database}); err != nil {
+		return nil, errors.Wrap(err, "PGDSN template is invalid")
+	}
+	config, err := pgx.ParseConfig(dsn.String())
+	if err != nil {
+		return nil, errors.Wrap(err, "rendered PGDSN is invalid")
 	}
 	return config, nil
 }