mirror of https://github.com/jackc/pgx.git
111 lines
2.2 KiB
Go
111 lines
2.2 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/jackc/pgconn"
|
|
"github.com/jackc/pgx"
|
|
"github.com/jackc/puddle"
|
|
)
|
|
|
|
type Pool struct {
|
|
p *puddle.Pool
|
|
}
|
|
|
|
// Connect creates a new Pool and immediately establishes one connection. ctx can be used to cancel this initial
|
|
// connection.
|
|
func Connect(ctx context.Context, connString string) (*Pool, error) {
|
|
p := &Pool{}
|
|
|
|
maxConnections := 5 // TODO - unhard-code
|
|
p.p = puddle.NewPool(
|
|
func(ctx context.Context) (interface{}, error) { return pgx.Connect(ctx, connString) },
|
|
func(value interface{}) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
value.(*pgx.Conn).Close(ctx)
|
|
cancel()
|
|
},
|
|
maxConnections)
|
|
|
|
// Initially establish one connection
|
|
res, err := p.p.Acquire(ctx)
|
|
if err != nil {
|
|
p.p.Close()
|
|
return nil, err
|
|
}
|
|
res.Release()
|
|
|
|
return p, nil
|
|
}
|
|
|
|
// Close closes all connections in the pool and rejects future Acquire calls. Blocks until all connections are returned
|
|
// to pool and closed.
|
|
func (p *Pool) Close() {
|
|
p.p.Close()
|
|
}
|
|
|
|
func (p *Pool) Acquire(ctx context.Context) (*Conn, error) {
|
|
res, err := p.p.Acquire(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Conn{res: res}, nil
|
|
}
|
|
|
|
func (p *Pool) Stat() *Stat {
|
|
return &Stat{s: p.p.Stat()}
|
|
}
|
|
|
|
func (p *Pool) Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error) {
|
|
c, err := p.Acquire(ctx)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer c.Release()
|
|
|
|
return c.Exec(ctx, sql, arguments...)
|
|
}
|
|
|
|
func (p *Pool) Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (*Rows, error) {
|
|
c, err := p.Acquire(ctx)
|
|
if err != nil {
|
|
return &Rows{err: err}, err
|
|
}
|
|
|
|
rows, err := c.Query(ctx, sql, optionsAndArgs...)
|
|
if err == nil {
|
|
rows.c = c
|
|
} else {
|
|
c.Release()
|
|
}
|
|
|
|
return rows, err
|
|
}
|
|
|
|
func (p *Pool) QueryRow(ctx context.Context, sql string, optionsAndArgs ...interface{}) *Row {
|
|
c, err := p.Acquire(ctx)
|
|
if err != nil {
|
|
return &Row{err: err}
|
|
}
|
|
|
|
row := c.QueryRow(ctx, sql, optionsAndArgs...)
|
|
row.c = c
|
|
return row
|
|
}
|
|
|
|
func (p *Pool) Begin() (*Tx, error) {
|
|
c, err := p.Acquire(context.Background())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
t, err := c.Begin()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Tx{t: t, c: c}, err
|
|
}
|