From 68e8bb8e43bfb8f27d0edb77ff2aac78c7eb1308 Mon Sep 17 00:00:00 2001 From: Richard Hayes Date: Tue, 20 Apr 2021 19:46:13 +0100 Subject: [PATCH 1/3] added new hook to Options for catering for dynamic passwords and the like --- .gitignore | 1 + base.go | 6 ++++++ db_test.go | 24 ++++++++++++++++++++++++ options.go | 4 ++++ 4 files changed, 35 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..9f11b755 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.idea/ diff --git a/base.go b/base.go index 985ea46d..27ae0f6a 100644 --- a/base.go +++ b/base.go @@ -115,6 +115,12 @@ func (db *baseDB) initConn(ctx context.Context, cn *pool.Conn) error { } } + if db.opt.BeforeConnect != nil { + if err := db.opt.BeforeConnect(ctx, db.opt); err != nil { + return err + } + } + err := db.startup(ctx, cn, db.opt.User, db.opt.Password, db.opt.Database, db.opt.ApplicationName) if err != nil { return err diff --git a/db_test.go b/db_test.go index d941a33b..8d9c512c 100644 --- a/db_test.go +++ b/db_test.go @@ -108,6 +108,30 @@ func TestDBConnectWithStartupNotice(t *testing.T) { require.NoError(t, db.Ping(context.Background()), "must successfully ping database with long application name") } +func TestBeforeConnect(t *testing.T) { + pwd := "dynamic-passwords-from-xkcd" + opt := pgOptions() + opt.BeforeConnect = func(ctx context.Context, o *pg.Options) error { + o.Password = pwd + return nil + } + + db := pg.Connect(opt) + defer db.Close() + + var val int + _, err := db.QueryOne(pg.Scan(&val), "SELECT 1") + if err != nil { + t.Fatal(err) + } + if val != 1 { + t.Fatalf(`got %q, wanted 1`, val) + } + if pwd != opt.Password { + t.Fatalf(`got %s, wanted %s`, opt.Password, pwd) + } +} + func TestOnConnect(t *testing.T) { opt := pgOptions() opt.OnConnect = func(ctx context.Context, db *pg.Conn) error { diff --git a/options.go b/options.go index 675041a9..2ef9aca1 100644 --- a/options.go +++ b/options.go @@ -32,6 +32,10 @@ type Options struct { // Network and Addr options. Dialer func(ctx context.Context, network, addr string) (net.Conn, error) + // BeforeConnect is a hook which is called before a new connection is + // established. Useful for scenarios where dynamic passwords are used + BeforeConnect func(ctx context.Context, opts *Options) error + // Hook that is called after new connection is established // and user is authenticated. OnConnect func(ctx context.Context, cn *Conn) error From 7b867afd99995efb6d3105aa1f32e84af0deef65 Mon Sep 17 00:00:00 2001 From: Richard Hayes Date: Wed, 21 Apr 2021 14:47:13 +0100 Subject: [PATCH 2/3] removed .gitignore --- .gitignore | 1 - 1 file changed, 1 deletion(-) delete mode 100644 .gitignore diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 9f11b755..00000000 --- a/.gitignore +++ /dev/null @@ -1 +0,0 @@ -.idea/ From 6186bafe8d93a81a544e07bfc2747a66ff6f1ed6 Mon Sep 17 00:00:00 2001 From: Richard Hayes Date: Fri, 29 Oct 2021 13:18:56 +0100 Subject: [PATCH 3/3] implemented options clone approach --- base.go | 17 ++++++++++------- options.go | 44 ++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/base.go b/base.go index 27ae0f6a..380cdc04 100644 --- a/base.go +++ b/base.go @@ -108,27 +108,30 @@ func (db *baseDB) initConn(ctx context.Context, cn *pool.Conn) error { } cn.Inited = true - if db.opt.TLSConfig != nil { - err := db.enableSSL(ctx, cn, db.opt.TLSConfig) + opt := db.opt.clone() + + if opt.TLSConfig != nil { + err := db.enableSSL(ctx, cn, opt.TLSConfig) if err != nil { return err } } - if db.opt.BeforeConnect != nil { - if err := db.opt.BeforeConnect(ctx, db.opt); err != nil { + + if opt.BeforeConnect != nil { + if err := opt.BeforeConnect(ctx, opt); err != nil { return err } } - err := db.startup(ctx, cn, db.opt.User, db.opt.Password, db.opt.Database, db.opt.ApplicationName) + err := db.startup(ctx, cn, opt.User, opt.Password, opt.Database, opt.ApplicationName) if err != nil { return err } - if db.opt.OnConnect != nil { + if opt.OnConnect != nil { p := pool.NewSingleConnPool(db.pool, cn) - return db.opt.OnConnect(ctx, newConn(ctx, db.withPool(p))) + return opt.OnConnect(ctx, newConn(ctx, db.withPool(p))) } return nil diff --git a/options.go b/options.go index 2ef9aca1..ddd76f40 100644 --- a/options.go +++ b/options.go @@ -11,6 +11,7 @@ import ( "runtime" "strconv" "strings" + "sync" "time" "go.opentelemetry.io/otel/attribute" @@ -20,6 +21,14 @@ import ( "github.com/go-pg/pg/v10/internal/pool" ) +type BeforeConnectOptions struct { + User string + Password string + + // TLS config for secure connections. + TLSConfig *tls.Config +} + // Options contains database connection options. type Options struct { // Network type, either tcp or unix. @@ -33,8 +42,9 @@ type Options struct { Dialer func(ctx context.Context, network, addr string) (net.Conn, error) // BeforeConnect is a hook which is called before a new connection is - // established. Useful for scenarios where dynamic passwords are used - BeforeConnect func(ctx context.Context, opts *Options) error + // established. Useful for scenarios where dynamic passwords are used. + // Timeout & Retry values set in this hook are not ignored + BeforeConnect func(ctx context.Context, o *Options) error // Hook that is called after new connection is established // and user is authenticated. @@ -98,6 +108,36 @@ type Options struct { // but idle connections are still discarded by the client // if IdleTimeout is set. IdleCheckFrequency time.Duration + + mux sync.Mutex +} + +func (opt *Options) clone() *Options { + return &Options{ + Network: opt.Network, + Addr: opt.Addr, + Dialer: opt.Dialer, + BeforeConnect: opt.BeforeConnect, + OnConnect: opt.OnConnect, + User: opt.User, + Password: opt.Password, + Database: opt.Database, + ApplicationName: opt.ApplicationName, + TLSConfig: opt.TLSConfig.Clone(), + DialTimeout: opt.DialTimeout, + ReadTimeout: opt.ReadTimeout, + WriteTimeout: opt.WriteTimeout, + MaxRetries: opt.MaxRetries, + RetryStatementTimeout: opt.RetryStatementTimeout, + MinRetryBackoff: opt.MinRetryBackoff, + MaxRetryBackoff: opt.MaxRetryBackoff, + PoolSize: opt.PoolSize, + MinIdleConns: opt.MinIdleConns, + MaxConnAge: opt.MaxConnAge, + PoolTimeout: opt.PoolTimeout, + IdleTimeout: opt.IdleTimeout, + IdleCheckFrequency: opt.IdleCheckFrequency, + } } func (opt *Options) init() {