db/postgres: support goormdb/postgres/v0.2.0

3 files changed, 56 insertions, 4 deletions

diff --git a/db/postgres/client.go b/db/postgres/client.go
index 49f8b2e..f8c6123 100644
--- a/db/postgres/client.go
+++ b/db/postgres/client.go
@@ -2,6 +2,7 @@ package postgres
 
 import (
         "context"
+        "database/sql"
         "net"
         "net/url"
         "strings"
@@ -9,9 +10,11 @@ import (
         "code.crute.us/mcrute/golib/secrets"
 
         "github.com/jackc/pgx/v5"
+        "github.com/jackc/pgx/v5/pgxpool"
+        "github.com/jackc/pgx/v5/stdlib"
 )
 
-func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, error) {
+func prepareUri(ctx context.Context, uri string, sc secrets.Client) (string, string, error) {
         // Prefix uri with postgres:// unless it already includes one of the
         // standard prefixes. Otherwise if scheme is omitted then url parsing
         // will fail to capture the username for secret lookup.
@@ -21,7 +24,7 @@ func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, err
 
         u, err := url.Parse(uri)
         if err != nil {
-                return nil, err
+                return "", "", err
         }
 
         // The username provided by the user (there should be no
@@ -31,7 +34,7 @@ func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, err
         // connect.
         cred, _, err := sc.DatabaseCredential(ctx, u.User.Username())
         if err != nil {
-                return nil, err
+                return "", "", err
         }
         u.User = url.UserPassword(cred.Username, cred.Password)
 
@@ -40,5 +43,50 @@ func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, err
                 u.Host = net.JoinHostPort(u.Host, "5432")
         }
 
-        return pgx.Connect(ctx, u.String())
+        return u.String(), u.Query().Get("TimeZone"), nil
+}
+
+func Connect(ctx context.Context, uri string, sc secrets.Client) (*pgx.Conn, error) {
+        dsn, tz, err := prepareUri(ctx, uri, sc)
+        if err != nil {
+                return nil, err
+        }
+
+        cfg, err := pgx.ParseConfig(dsn)
+        if err != nil {
+                return nil, err
+        }
+
+        if tz != "" {
+                cfg.RuntimeParams["timezone"] = tz
+        }
+
+        return pgx.ConnectConfig(ctx, cfg)
+}
+
+func ConnectPool(ctx context.Context, uri string, sc secrets.Client) (*pgxpool.Pool, error) {
+        dsn, tz, err := prepareUri(ctx, uri, sc)
+        if err != nil {
+                return nil, err
+        }
+
+        cfg, err := pgxpool.ParseConfig(dsn)
+        if err != nil {
+                return nil, err
+        }
+
+        if tz != "" {
+                cfg.ConnConfig.RuntimeParams["timezone"] = tz
+        }
+
+        return pgxpool.NewWithConfig(ctx, cfg)
+}
+
+func ConnectStdlibPool(ctx context.Context, uri string, sc secrets.Client) (*sql.DB, error) {
+        pool, err := ConnectPool(ctx, uri, sc)
+        if err != nil {
+                return nil, err
+        }
+
+        return stdlib.OpenDBFromPool(pool), nil
 }
diff --git a/db/postgres/go.mod b/db/postgres/go.mod
index ae0e3d0..ac3bee8 100644
--- a/db/postgres/go.mod
+++ b/db/postgres/go.mod
@@ -37,6 +37,7 @@ require (
         github.com/hashicorp/yamux v0.0.0-20180604194846-3520598351bb // indirect
         github.com/jackc/pgpassfile v1.0.0 // indirect
         github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+        github.com/jackc/puddle/v2 v2.2.1 // indirect
         github.com/mattn/go-colorable v0.1.6 // indirect
         github.com/mattn/go-isatty v0.0.12 // indirect
         github.com/mitchellh/copystructure v1.0.0 // indirect
@@ -51,6 +52,7 @@ require (
         go.uber.org/atomic v1.9.0 // indirect
         golang.org/x/crypto v0.9.0 // indirect
         golang.org/x/net v0.10.0 // indirect
+        golang.org/x/sync v0.1.0 // indirect
         golang.org/x/sys v0.8.0 // indirect
         golang.org/x/text v0.9.0 // indirect
         golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 // indirect
diff --git a/db/postgres/go.sum b/db/postgres/go.sum
index 78d421b..950a9b7 100644
--- a/db/postgres/go.sum
+++ b/db/postgres/go.sum
@@ -152,6 +152,7 @@ github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZ
 github.com/jackc/pgx/v5 v5.5.0 h1:NxstgwndsTRy7eq9/kqYc/BZh5w2hHJV86wjvO+1xPw=
 github.com/jackc/pgx/v5 v5.5.0/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
 github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
+github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
 github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
 github.com/jhump/protoreflect v1.6.0 h1:h5jfMVslIg6l29nsMs0D8Wj17RDVdNYti0vDN/PZZoE=
 github.com/jhump/protoreflect v1.6.0/go.mod h1:eaTn3RZAmMBcV0fifFvlm6VHNz3wSkYyXYWUh7ymB74=
@@ -273,6 +274,7 @@ golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=