From b252e2d68a8655e9acdf3457046754f4d8fd3913 Mon Sep 17 00:00:00 2001 From: Robert Lin Date: Thu, 9 May 2024 14:42:06 -0700 Subject: [PATCH] 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 --- .../cloudsql/BUILD.bazel | 14 +++ .../cloudsql/cloudsql.go | 93 +++++++++++++++++++ .../runtime/contract/BUILD.bazel | 1 + .../runtime/contract/postgresql.go | 76 ++++++++------- 4 files changed, 150 insertions(+), 34 deletions(-) create mode 100644 lib/managedservicesplatform/cloudsql/BUILD.bazel create mode 100644 lib/managedservicesplatform/cloudsql/cloudsql.go diff --git a/lib/managedservicesplatform/cloudsql/BUILD.bazel b/lib/managedservicesplatform/cloudsql/BUILD.bazel new file mode 100644 index 00000000000..8fef0bd9ee8 --- /dev/null +++ b/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", + ], +) diff --git a/lib/managedservicesplatform/cloudsql/cloudsql.go b/lib/managedservicesplatform/cloudsql/cloudsql.go new file mode 100644 index 00000000000..18357278261 --- /dev/null +++ b/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 +} diff --git a/lib/managedservicesplatform/runtime/contract/BUILD.bazel b/lib/managedservicesplatform/runtime/contract/BUILD.bazel index 717c6aa0ece..5c756c832eb 100644 --- a/lib/managedservicesplatform/runtime/contract/BUILD.bazel +++ b/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", diff --git a/lib/managedservicesplatform/runtime/contract/postgresql.go b/lib/managedservicesplatform/runtime/contract/postgresql.go index 019ccf83217..fad918f857e 100644 --- a/lib/managedservicesplatform/runtime/contract/postgresql.go +++ b/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 }