Apache Beam Go连接Cloud SQL PostgreSQL报错：未知驱动cloudsql-postgres

问题根源 #

你遇到的sql: unknown driver "cloudsql-postgres"错误，核心原因是Dataflow的分布式执行模型导致驱动注册只在主进程生效，worker进程无法获取到注册的驱动：

• 你在ReadFromPostgres函数里调用pgxv4.RegisterDriver注册驱动，这个操作仅在Dataflow的主（Driver）进程中执行。
• Dataflow的作业逻辑（包括databaseio.Query的queryFn）是在独立的Worker进程中运行的，Worker进程不会继承主进程中注册的驱动。
• 你用defer cleanup()会在主进程执行完ReadFromPostgres后立即清理已注册的驱动，进一步导致Worker进程启动时驱动已被移除。

修复方案 #

方案1：在init()函数中注册驱动（推荐） #

把驱动注册逻辑放到全局init()函数中，确保Worker进程启动时自动完成驱动注册，且不要调用cleanup（Worker进程需要全程使用驱动）：

import (
    "context"
    "fmt"
    "reflect"
    "time"

    "cloud.google.com/go/bigquery"
    "cloud.google.com/go/cloudsqlconn"
    "cloud.google.com/go/cloudsqlconn/postgres/pgxv4"
    "github.com/apache/beam/sdks/v2/go/pkg/beam"
    "github.com/apache/beam/sdks/v2/go/pkg/beam/io/bigqueryio"
    "github.com/apache/beam/sdks/v2/go/pkg/beam/io/databaseio"
    "google.golang.org/api/iterator"
)

func init() {
    // 注册cloudsql-postgres驱动，不执行cleanup
    _, err := pgxv4.RegisterDriver("cloudsql-postgres", cloudsqlconn.WithIAMAuthN())
    if err != nil {
        panic(fmt.Sprintf("failed to register cloudsql-postgres driver: %v", err))
    }
}

func ReadFromPostgres(
    ctx context.Context,
    scope beam.Scope,
    opt *SourceOption,
    secretReader *SecretReader,
    elemType reflect.Type,
    query string,
    postgresDriver string,
) (beam.PCollection, error) {
    scope = scope.Scope("Read from database")
    secret, err := opt.Database.DSN.GetValue(ctx, secretReader)
    if err != nil {
        return beam.PCollection{}, fmt.Errorf("error getting secret value: %w", err)
    }
    envName, location, dbInstanceName, user, dbName := opt.Database.PGEnv.EnvName, opt.Database.PGEnv.Location, opt.Database.PGEnv.DBInstanceName, opt.Database.PGEnv.User, opt.Database.PGEnv.DBName
    // 修正sslmode为符合IAM认证的要求值
    dsn := fmt.Sprintf("host=%s:%s:%s user=%s dbname=%s password=%s sslmode=require", envName, location, dbInstanceName, user, dbName, secret)
    return databaseio.Query(scope, postgresDriver, dsn, query, elemType), nil
}

方案2：自定义QueryFn在Worker侧注册驱动 #

如果不想用全局init()，可以自定义databaseio.Query的底层函数，在Worker进程执行时注册驱动：

import (
    // 保留原有导入
    "cloud.google.com/go/cloudsqlconn"
    "cloud.google.com/go/cloudsqlconn/postgres/pgxv4"
    "github.com/apache/beam/sdks/v2/go/pkg/beam/io/databaseio"
)

// 自定义QueryFn，在Open方法中注册驱动
type customQueryFn struct {
    databaseio.QueryFn
    DriverName string
}

func (fn *customQueryFn) Open(ctx context.Context) (databaseio.Reader, error) {
    // 在Worker进程中注册驱动
    _, err := pgxv4.RegisterDriver(fn.DriverName, cloudsqlconn.WithIAMAuthN())
    if err != nil {
        return nil, fmt.Errorf("failed to register driver: %w", err)
    }
    // 调用原有的Open逻辑
    return fn.QueryFn.Open(ctx)
}

func ReadFromPostgres(
    ctx context.Context,
    scope beam.Scope,
    opt *SourceOption,
    secretReader *SecretReader,
    elemType reflect.Type,
    query string,
    postgresDriver string,
) (beam.PCollection, error) {
    scope = scope.Scope("Read from database")
    secret, err := opt.Database.DSN.GetValue(ctx, secretReader)
    if err != nil {
        return beam.PCollection{}, fmt.Errorf("error getting secret value: %w", err)
    }
    envName, location, dbInstanceName, user, dbName := opt.Database.PGEnv.EnvName, opt.Database.PGEnv.Location, opt.Database.PGEnv.DBInstanceName, opt.Database.PGEnv.User, opt.Database.PGEnv.DBName
    dsn := fmt.Sprintf("host=%s:%s:%s user=%s dbname=%s password=%s sslmode=require", envName, location, dbInstanceName, user, dbName, secret)
    
    // 创建自定义QueryFn并执行
    qfn := &customQueryFn{
        QueryFn: databaseio.QueryFn{
            DriverName: postgresDriver,
            DSN:        dsn,
            Query:      query,
            ElemType:   elemType,
        },
        DriverName: postgresDriver,
    }
    return beam.ParDo(scope, qfn, beam.Impulse(scope)), nil
}

额外注意事项 #

• 你的DSN中sslmode=disable不符合Cloud SQL的IAM认证要求，建议改为sslmode=require（Cloud SQL会自动提供SSL证书，无需手动配置）。
• 确保Dataflow Worker的服务账号拥有cloudsql.client权限，以及Cloud SQL实例的IAM认证权限（如果使用IAM AuthN）。

内容的提问来源于stack exchange，提问作者snark17