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