1// +build go1.10
2
3package stdlib
4
5import (
6	"context"
7	"database/sql"
8	"database/sql/driver"
9
10	"github.com/jackc/pgx"
11)
12
13// OptionOpenDB options for configuring the driver when opening a new db pool.
14type OptionOpenDB func(*connector)
15
16// OptionAfterConnect provide a callback for after connect.
17func OptionAfterConnect(ac func(*pgx.Conn) error) OptionOpenDB {
18	return func(dc *connector) {
19		dc.AfterConnect = ac
20	}
21}
22
23func OpenDB(config pgx.ConnConfig, opts ...OptionOpenDB) *sql.DB {
24	c := connector{
25		ConnConfig:   config,
26		AfterConnect: func(*pgx.Conn) error { return nil }, // noop after connect by default
27		driver:       pgxDriver,
28	}
29
30	for _, opt := range opts {
31		opt(&c)
32	}
33
34	return sql.OpenDB(c)
35}
36
37type connector struct {
38	pgx.ConnConfig
39	AfterConnect func(*pgx.Conn) error // function to call on every new connection
40	driver       *Driver
41}
42
43// Connect implement driver.Connector interface
44func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
45	var (
46		err  error
47		conn *pgx.Conn
48	)
49
50	if conn, err = pgx.Connect(c.ConnConfig); err != nil {
51		return nil, err
52	}
53
54	if err = c.AfterConnect(conn); err != nil {
55		return nil, err
56	}
57
58	return &Conn{conn: conn, driver: c.driver, connConfig: c.ConnConfig}, nil
59}
60
61// Driver implement driver.Connector interface
62func (c connector) Driver() driver.Driver {
63	return c.driver
64}
65