Skip to content

BeforeConnect Hook #1875

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: v10
Choose a base branch
from
Open

BeforeConnect Hook #1875

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
68e8bb8
added new hook to Options for catering for dynamic passwords and the …
Apr 20, 2021
7b867af
removed .gitignore
Apr 21, 2021
98f3ff4
Merge branch 'go-pg:v10' into v10
justcompile Oct 29, 2021
6186baf
implemented options clone approach
Oct 29, 2021
22a4c8d
Merge branch 'v10' of github.com:go-pg/pg into v10
Oct 29, 2021
5a91137
Merge branch 'v10' of github.com:justcompile/pg into v10
Oct 29, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Copy link
Member

@vmihailenco vmihailenco Apr 21, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe we should copy the options to avoid data race (since this hooks is called concurrently):

opt := db.opt.Clone()
// use the cloned opt

Copy link
Author

@justcompile justcompile Apr 21, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and then reassign to db.opt?

if db.opt.BeforeConnect != nil {
    opts := db.opt.Clone()
    if err := opt.BeforeConnect(ctx, opt); err != nil { return err }
    db.opts = opts		
}

Copy link
Collaborator

@elliotcourant elliotcourant Apr 21, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you re assign to db.opt then you need to make that re assignment atomic (*thread safe?). What if this code is being executed in multiple go routines?

Copy link
Collaborator

@elliotcourant elliotcourant Apr 21, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also if we replace the db.opt here this can lead to a few interesting issues depending on how the user is using the library. What if they change the DB address entirely? what if they add TLS? Beyond simply changing the password or the user that a connection is issued with; changing other variables available in the options object could have a lot of unintended consequences that are currently unknown.

If at all possible we might just want to limit being able to change authentication parameters if we do update the original db.opts. Maybe limiting it to changing the user, password and TLS config.

If they for some reason changed something like the connection pool settings, how would that affect the pool that is already running?

Copy link
Member

@vmihailenco vmihailenco Apr 23, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@justcompile I don't think you need to replace db.opt. Just clone it (when needed for BeforeConnect) and use below in startup:

opt := db.opt.Clone()
opt.BeforeConnect(ctx, opt)
db.startup(ctx, cn, opt.User, opt.Password, opt.Database, opt.ApplicationName)

That will not cover all options but is good enough for this particular use case. We can improve it later if needed.

Copy link
Author

@justcompile justcompile Apr 23, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

if db.opt.BeforeConnect != nil {
	db.opt.mux.Lock()
	updateOpts := &BeforeConnectOptions{
		User: db.opt.User,
		Password: db.opt.Password,
		TLSConfig: db.opt.TLSConfig,
	}
	if err := db.opt.BeforeConnect(ctx, updateOpts); err != nil {
		db.opt.mux.Unlock()
		return err
	}
		
	db.opt.applyUpdatedConnectionOptions(updateOpts)
	db.opt.mux.Unlock()
}
	
// options.go
func (opt *Options) applyUpdatedConnectionOptions(bcOpts *BeforeConnectOptions) {
	opt.User = bcOpts.User
	opt.Password = bcOpts.Password

	if bcOpts.TLSConfig != nil {
		opt.TLSConfig = bcOpts.TLSConfig.Clone()
	}
}

Based on the concurrency issue mentioned & the points around unintended side effects of a user changing the Addr or Pool options, I was looking at something like the above.

But if that's overkill, I can go with the opts.Clone approach.

I'm also locking the mux here rather than within applyUpdatedConnectionOptions so that if multiple go routines are calling BeforeConnect it's a bit more deterministic

Copy link
Member

@vmihailenco vmihailenco Apr 23, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's keep as is BeforeConnect(ctx context.Context, o *pg.Options) error. It will not work for all options but that is fine. It is an advanced feature. Add add a warning to the comment.

return err
}
}

err := db.startup(ctx, cn, db.opt.User, db.opt.Password, db.opt.Database, db.opt.ApplicationName)
if err != nil {
return err
Expand Down
24 changes: 24 additions & 0 deletions db_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Copy link
Collaborator

@elliotcourant elliotcourant Apr 21, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add another test that tests the before connect with multiple go routines to make sure that a race condition is not triggered by accident. A race condition in a before connection hook for an ORM could cripple an application.

Copy link
Collaborator

@elliotcourant elliotcourant Oct 30, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a good way to do this would be to set the options to have a max pool of 1. And then try to send multiple queries concurrently in a few go routines.

This should trigger a concurrent call of the BeforeConnect method; and given the -race flag should be a sufficient enough smoke test for most use cases.

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 {
Expand Down
4 changes: 4 additions & 0 deletions options.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down