143 lines
3.5 KiB
Go
143 lines
3.5 KiB
Go
|
package pg
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"fmt"
|
||
|
"time"
|
||
|
|
||
|
"github.com/go-pg/pg/v10/internal/pool"
|
||
|
"github.com/go-pg/pg/v10/orm"
|
||
|
)
|
||
|
|
||
|
// Connect connects to a database using provided options.
|
||
|
//
|
||
|
// The returned DB is safe for concurrent use by multiple goroutines
|
||
|
// and maintains its own connection pool.
|
||
|
func Connect(opt *Options) *DB {
|
||
|
opt.init()
|
||
|
return newDB(
|
||
|
context.Background(),
|
||
|
&baseDB{
|
||
|
opt: opt,
|
||
|
pool: newConnPool(opt),
|
||
|
fmter: orm.NewFormatter(),
|
||
|
},
|
||
|
)
|
||
|
}
|
||
|
|
||
|
func newDB(ctx context.Context, baseDB *baseDB) *DB {
|
||
|
db := &DB{
|
||
|
baseDB: baseDB.clone(),
|
||
|
ctx: ctx,
|
||
|
}
|
||
|
db.baseDB.db = db
|
||
|
return db
|
||
|
}
|
||
|
|
||
|
// DB is a database handle representing a pool of zero or more
|
||
|
// underlying connections. It's safe for concurrent use by multiple
|
||
|
// goroutines.
|
||
|
type DB struct {
|
||
|
*baseDB
|
||
|
ctx context.Context
|
||
|
}
|
||
|
|
||
|
var _ orm.DB = (*DB)(nil)
|
||
|
|
||
|
func (db *DB) String() string {
|
||
|
return fmt.Sprintf("DB<Addr=%q%s>", db.opt.Addr, db.fmter)
|
||
|
}
|
||
|
|
||
|
// Options returns read-only Options that were used to connect to the DB.
|
||
|
func (db *DB) Options() *Options {
|
||
|
return db.opt
|
||
|
}
|
||
|
|
||
|
// Context returns DB context.
|
||
|
func (db *DB) Context() context.Context {
|
||
|
return db.ctx
|
||
|
}
|
||
|
|
||
|
// WithContext returns a copy of the DB that uses the ctx.
|
||
|
func (db *DB) WithContext(ctx context.Context) *DB {
|
||
|
return newDB(ctx, db.baseDB)
|
||
|
}
|
||
|
|
||
|
// WithTimeout returns a copy of the DB that uses d as the read/write timeout.
|
||
|
func (db *DB) WithTimeout(d time.Duration) *DB {
|
||
|
return newDB(db.ctx, db.baseDB.WithTimeout(d))
|
||
|
}
|
||
|
|
||
|
// WithParam returns a copy of the DB that replaces the param with the value
|
||
|
// in queries.
|
||
|
func (db *DB) WithParam(param string, value interface{}) *DB {
|
||
|
return newDB(db.ctx, db.baseDB.WithParam(param, value))
|
||
|
}
|
||
|
|
||
|
// Listen listens for notifications sent with NOTIFY command.
|
||
|
func (db *DB) Listen(ctx context.Context, channels ...string) *Listener {
|
||
|
ln := &Listener{
|
||
|
db: db,
|
||
|
}
|
||
|
ln.init()
|
||
|
_ = ln.Listen(ctx, channels...)
|
||
|
return ln
|
||
|
}
|
||
|
|
||
|
// Conn represents a single database connection rather than a pool of database
|
||
|
// connections. Prefer running queries from DB unless there is a specific
|
||
|
// need for a continuous single database connection.
|
||
|
//
|
||
|
// A Conn must call Close to return the connection to the database pool
|
||
|
// and may do so concurrently with a running query.
|
||
|
//
|
||
|
// After a call to Close, all operations on the connection fail.
|
||
|
type Conn struct {
|
||
|
*baseDB
|
||
|
ctx context.Context
|
||
|
}
|
||
|
|
||
|
var _ orm.DB = (*Conn)(nil)
|
||
|
|
||
|
// Conn returns a single connection from the connection pool.
|
||
|
// Queries run on the same Conn will be run in the same database session.
|
||
|
//
|
||
|
// Every Conn must be returned to the database pool after use by
|
||
|
// calling Conn.Close.
|
||
|
func (db *DB) Conn() *Conn {
|
||
|
return newConn(db.ctx, db.baseDB.withPool(pool.NewStickyConnPool(db.pool)))
|
||
|
}
|
||
|
|
||
|
func newConn(ctx context.Context, baseDB *baseDB) *Conn {
|
||
|
conn := &Conn{
|
||
|
baseDB: baseDB,
|
||
|
ctx: ctx,
|
||
|
}
|
||
|
conn.baseDB.db = conn
|
||
|
return conn
|
||
|
}
|
||
|
|
||
|
// Context returns DB context.
|
||
|
func (db *Conn) Context() context.Context {
|
||
|
if db.ctx != nil {
|
||
|
return db.ctx
|
||
|
}
|
||
|
return context.Background()
|
||
|
}
|
||
|
|
||
|
// WithContext returns a copy of the DB that uses the ctx.
|
||
|
func (db *Conn) WithContext(ctx context.Context) *Conn {
|
||
|
return newConn(ctx, db.baseDB)
|
||
|
}
|
||
|
|
||
|
// WithTimeout returns a copy of the DB that uses d as the read/write timeout.
|
||
|
func (db *Conn) WithTimeout(d time.Duration) *Conn {
|
||
|
return newConn(db.ctx, db.baseDB.WithTimeout(d))
|
||
|
}
|
||
|
|
||
|
// WithParam returns a copy of the DB that replaces the param with the value
|
||
|
// in queries.
|
||
|
func (db *Conn) WithParam(param string, value interface{}) *Conn {
|
||
|
return newConn(db.ctx, db.baseDB.WithParam(param, value))
|
||
|
}
|