Skip to content

Commit

Permalink
feat: add pgx v5 support (#639) (#642)
Browse files Browse the repository at this point in the history
* chore: bump pgx version to v5 (#639)
  • Loading branch information
annafang-google committed Oct 12, 2023
1 parent bd821fb commit 8d86d92
Show file tree
Hide file tree
Showing 5 changed files with 181 additions and 66 deletions.
106 changes: 84 additions & 22 deletions e2e_postgres_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,12 @@ import (
"time"

"cloud.google.com/go/cloudsqlconn"
"github.com/jackc/pgx/v4"
"github.com/jackc/pgx/v5"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"

"cloud.google.com/go/cloudsqlconn/postgres/pgxv4"
"cloud.google.com/go/cloudsqlconn/postgres/pgxv5"
)

var (
Expand Down Expand Up @@ -151,7 +152,26 @@ func TestEngineVersion(t *testing.T) {
}
}

func TestPostgresHook(t *testing.T) {
func TestPostgresV5Hook(t *testing.T) {
tests := []struct {
driver string
source string
IAMAuthN bool
}{
{
driver: "cloudsql-postgres",
source: fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=disable",
postgresConnName, postgresUser, postgresPass, postgresDB),
IAMAuthN: false,
},
{
driver: "cloudsql-postgres-iam",
source: fmt.Sprintf("host=%s user=%s dbname=%s sslmode=disable",
postgresConnName, postgresUserIAM, postgresDB),
IAMAuthN: true,
},
}

if testing.Short() {
t.Skip("skipping Postgres integration tests")
}
Expand All @@ -162,29 +182,71 @@ func TestPostgresHook(t *testing.T) {
}
t.Log(now)
}
pgxv4.RegisterDriver("cloudsql-postgres")
db, err := sql.Open(
"cloudsql-postgres",
fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=disable",
postgresConnName, postgresUser, postgresPass, postgresDB),
)
if err != nil {
t.Fatalf("sql.Open want err = nil, got = %v", err)

for _, tc := range tests {
if tc.IAMAuthN {
pgxv5.RegisterDriver(tc.driver, cloudsqlconn.WithIAMAuthN())
} else {
pgxv5.RegisterDriver(tc.driver)
}
db, err := sql.Open(tc.driver, tc.source)

if err != nil {
t.Fatalf("sql.Open want err = nil, got = %v", err)
}
defer db.Close()
testConn(db)
}
defer db.Close()
testConn(db)
}

pgxv4.RegisterDriver("cloudsql-postgres-iam", cloudsqlconn.WithIAMAuthN())
db2, err := sql.Open(
"cloudsql-postgres-iam",
fmt.Sprintf("host=%s user=%s dbname=%s sslmode=disable",
postgresConnName, postgresUserIAM, postgresDB),
)
if err != nil {
t.Fatalf("sql.Open want err = nil, got = %v", err)
func TestPostgresV4Hook(t *testing.T) {
tests := []struct {
driver string
source string
IAMAuthN bool
}{
{
driver: "cloudsql-postgres",
source: fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=disable",
postgresConnName, postgresUser, postgresPass, postgresDB),
IAMAuthN: false,
},
{
driver: "cloudsql-postgres-iam",
source: fmt.Sprintf("host=%s user=%s dbname=%s sslmode=disable",
postgresConnName, postgresUserIAM, postgresDB),
IAMAuthN: true,
},
}
if testing.Short() {
t.Skip("skipping Postgres integration tests")
}
testConn := func(db *sql.DB) {
var now time.Time
if err := db.QueryRow("SELECT NOW()").Scan(&now); err != nil {
t.Fatalf("QueryRow failed: %v", err)
}
t.Log(now)
}

for _, tc := range tests {
defer func() {
if err := recover(); err != nil {
t.Log("cloudsql-postgres register panic occurred:", err)
}
}()
if tc.IAMAuthN {
pgxv4.RegisterDriver(tc.driver, cloudsqlconn.WithIAMAuthN())
} else {
pgxv4.RegisterDriver(tc.driver)
}
db, err := sql.Open(tc.driver, tc.source)
if err != nil {
t.Fatalf("sql.Open want err = nil, got = %v", err)
}
defer db.Close()
testConn(db)
}
defer db2.Close()
testConn(db2)
}

// removeAuthEnvVar retrieves an OAuth2 token and a path to a service account key
Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ require (
github.com/go-sql-driver/mysql v1.7.1
github.com/google/uuid v1.3.1
github.com/jackc/pgx/v4 v4.18.1
github.com/jackc/pgx/v5 v5.4.3
github.com/microsoft/go-mssqldb v1.6.0
go.opencensus.io v0.24.0
golang.org/x/net v0.17.0
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,8 @@ github.com/jackc/pgx/v4 v4.0.0-pre1.0.20190824185557-6972a5742186/go.mod h1:X+GQ
github.com/jackc/pgx/v4 v4.12.1-0.20210724153913-640aa07df17c/go.mod h1:1QD0+tgSXP7iUjYm9C1NxKhny7lq6ee99u/z+IHFcgs=
github.com/jackc/pgx/v4 v4.18.1 h1:YP7G1KABtKpB5IHrO9vYwSrCOhs7p3uqhvhhQBptya0=
github.com/jackc/pgx/v4 v4.18.1/go.mod h1:FydWkUyadDmdNH/mHnGob881GawxeEm7TcMCzkb+qQE=
github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
Expand Down
44 changes: 0 additions & 44 deletions postgres/pgxv4/postgres_test.go

This file was deleted.

94 changes: 94 additions & 0 deletions postgres/pgxv5/postgres.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// 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 pgxv5 provides a Cloud SQL Postgres driver that uses pgx v5 and works
// with the database/sql package.
package pgxv5

import (
"context"
"database/sql"
"database/sql/driver"
"net"
"sync"

"cloud.google.com/go/cloudsqlconn"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/stdlib"
)

// RegisterDriver registers a Postgres driver that uses the cloudsqlconn.Dialer
// configured with the provided options. The choice of name is entirely up to
// the caller and may be used to distinguish between multiple registrations of
// differently configured Dialers. The driver uses pgx/v5 internally.
// RegisterDriver returns a cleanup function that should be called one the
// database connection is no longer needed.
func RegisterDriver(name string, opts ...cloudsqlconn.Option) (func() error, error) {
d, err := cloudsqlconn.NewDialer(context.Background(), opts...)
if err != nil {
return func() error { return nil }, err
}
sql.Register(name, &pgDriver{
d: d,
dbURIs: make(map[string]string),
})
return func() error { return d.Close() }, nil
}

type pgDriver struct {
d *cloudsqlconn.Dialer
mu sync.RWMutex
// dbURIs is a map of DSN to DB URI for registered connection names.
dbURIs map[string]string
}

// Open accepts a keyword/value formatted connection string and returns a
// connection to the database using cloudsqlconn.Dialer. The Cloud SQL instance
// connection name should be specified in the host field. For example:
//
// "host=my-project:us-central1:my-db-instance user=myuser password=mypass"
func (p *pgDriver) Open(name string) (driver.Conn, error) {
dbURI, err := p.dbURI(name)
if err != nil {
return nil, err
}
return stdlib.GetDefaultDriver().Open(dbURI)

}

// dbURI registers a driver using the provided DSN. If the name has already
// been registered, dbURI returns the existing registration.
func (p *pgDriver) dbURI(name string) (string, error) {
p.mu.Lock()
defer p.mu.Unlock()
dbURI, ok := p.dbURIs[name]
if ok {
return dbURI, nil
}

config, err := pgx.ParseConfig(name)
if err != nil {
return "", err
}
instConnName := config.Config.Host // Extract instance connection name
config.Config.Host = "localhost" // Replace it with a default value
config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
return p.d.Dial(ctx, instConnName)
}

dbURI = stdlib.RegisterConnConfig(config)
p.dbURIs[name] = dbURI

return dbURI, nil
}

0 comments on commit 8d86d92

Please sign in to comment.