mirror of
1
Fork 0
gotosocial/vendor/github.com/go-pg/pg/v10/db.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))
}