func (s *S) SetUpSuite(c *C) {
s.discoverd, s.cleanup = setup(c)
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxConfig := newPgxConnPoolConfig()
pgxpool, err := pgx.NewConnPool(pgxConfig)
if err != nil {
c.Fatal(err)
}
db := postgres.New(pgxpool, nil)
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
db.Close()
// reconnect with prepared statements
pgxConfig.AfterConnect = schema.PrepareStatements
pgxpool, err = pgx.NewConnPool(pgxConfig)
if err != nil {
c.Fatal(err)
}
db = postgres.New(pgxpool, nil)
s.pgx = db.ConnPool
s.pgx.Exec(sqlCreateTruncateTables)
}
func TestNewConnPool(t *testing.T) {
t.Parallel()
var numCallbacks int
afterConnect := func(c *pgx.Conn) error {
numCallbacks++
return nil
}
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 2, AfterConnect: afterConnect}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatal("Unable to establish connection pool")
}
defer pool.Close()
// It initially connects once
stat := pool.Stat()
if stat.CurrentConnections != 1 {
t.Errorf("Expected 1 connection to be established immediately, but %v were", numCallbacks)
}
// Pool creation returns an error if any AfterConnect callback does
errAfterConnect := errors.New("Some error")
afterConnect = func(c *pgx.Conn) error {
return errAfterConnect
}
config = pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 2, AfterConnect: afterConnect}
pool, err = pgx.NewConnPool(config)
if err != errAfterConnect {
t.Errorf("Expected errAfterConnect but received unexpected: %v", err)
}
}
func TestAcquireTimeoutSanity(t *testing.T) {
t.Parallel()
config := pgx.ConnPoolConfig{
ConnConfig: *defaultConnConfig,
MaxConnections: 1,
}
// case 1: default 0 value
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Expected NewConnPool with default config.AcquireTimeout not to fail, instead it failed with '%v'", err)
}
pool.Close()
// case 2: negative value
config.AcquireTimeout = -1 * time.Second
_, err = pgx.NewConnPool(config)
if err == nil {
t.Fatal("Expected NewConnPool with negative config.AcquireTimeout to fail, instead it did not")
}
// case 3: positive value
config.AcquireTimeout = 1 * time.Second
pool, err = pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Expected NewConnPool with positive config.AcquireTimeout not to fail, instead it failed with '%v'", err)
}
defer pool.Close()
}
func (c *ConnectionInfo) Setup() {
var err interface{}
c.Pool, err = pgx.NewConnPool(getConfig())
if err != nil {
log.Fatalf("Initializing PostgreSQL connection pool failed: %v", err)
}
}
func createConnPool() (*pgx.ConnPool, error) {
var config pgx.ConnPoolConfig
var err error
config.ConnConfig, err = pgx.ParseEnvLibpq()
if err != nil {
return nil, err
}
if config.Host == "" {
config.Host = "localhost"
}
if config.User == "" {
config.User = os.Getenv("USER")
}
if config.Database == "" {
config.Database = "pgxdata"
}
config.TLSConfig = nil
config.UseFallbackTLS = false
config.MaxConnections = 10
return pgx.NewConnPool(config)
}
func (d *DB) Open() error {
conf := pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: d.config.Host,
Port: uint16(d.config.Port),
User: d.config.Username,
Password: d.config.Password,
Database: d.config.Database,
Logger: d.config.Logger,
},
MaxConnections: d.config.MaxConnections,
}
p, err := pgx.NewConnPool(conf)
if err != nil {
return err
}
db, err := stdlib.OpenFromConnPool(p)
if err != nil {
return err
}
dbx := sqlx.NewDb(db, "pgx")
d.queryer = queryer{impl: dbx}
return nil
}
func TestPoolWithAcquireTimeoutSet(t *testing.T) {
t.Parallel()
connAllocTimeout := 2 * time.Second
config := pgx.ConnPoolConfig{
ConnConfig: *defaultConnConfig,
MaxConnections: 1,
AcquireTimeout: connAllocTimeout,
}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
// Consume all connections ...
allConnections := acquireAllConnections(t, pool, config.MaxConnections)
defer releaseAllConnections(pool, allConnections)
// ... then try to consume 1 more. It should fail after a short timeout.
_, timeTaken, err := acquireWithTimeTaken(pool)
if err == nil || err.Error() != "Timeout: All connections in pool are busy" {
t.Fatalf("Expected error to be 'Timeout: All connections in pool are busy', instead it was '%v'", err)
}
if timeTaken < connAllocTimeout {
t.Fatalf("Expected connection allocation time to be at least %v, instead it was '%v'", connAllocTimeout, timeTaken)
}
}
func OpenDatabaseConnection(config configuration.DatabaseConfiguration) (pool *pgx.ConnPool, err error) {
connectionUri := fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=%s",
config.User,
config.Password,
config.Host,
config.Port,
config.Database,
config.SSLMode)
connectionConfig, err := pgx.ParseURI(connectionUri)
if err != nil {
return pool, err
}
maxConnections := 50
poolConfig := pgx.ConnPoolConfig{connectionConfig, maxConnections, nil}
pool, err = pgx.NewConnPool(poolConfig)
if err != nil {
return pool, err
}
Database = pool
return pool, err
}
func TestOpenFromConnPool(t *testing.T) {
connConfig := pgx.ConnConfig{
Host: "127.0.0.1",
User: "******",
Password: "******",
Database: "pgx_test",
}
config := pgx.ConnPoolConfig{ConnConfig: connConfig}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
db, err := stdlib.OpenFromConnPool(pool)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer closeDB(t, db)
// Can get pgx.ConnPool from driver
driver := db.Driver().(*stdlib.Driver)
if driver.Pool == nil {
t.Fatal("Expected driver opened through OpenFromConnPool to have Pool, but it did not")
}
// Normal sql/database still works
var n int64
err = db.QueryRow("select 1").Scan(&n)
if err != nil {
t.Fatalf("db.QueryRow unexpectedly failed: %v", err)
}
}
// Connect connects to the database using env vars.
// After connect, it creates tables if missing.
func (p *DB) Connect() (err error) {
cfg, err := pgx.ParseEnvLibpq()
if err != nil {
return
}
pool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: cfg,
MaxConnections: 25,
})
if err != nil {
return
}
c, err := pgx_stdlib.OpenFromConnPool(pool)
if err != nil {
return
}
p.conn, err = gorm.Open("postgres", c)
if err != nil {
return
}
d := p.conn.AutoMigrate(&secrets.Secret{}, &secrets.Key{})
return d.Error
}
func main() {
var err error
connPoolConfig := pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "127.0.0.1",
User: "******",
Password: "******",
Database: "url_shortener",
Logger: log.New("module", "pgx"),
},
MaxConnections: 5,
AfterConnect: afterConnect,
}
pool, err = pgx.NewConnPool(connPoolConfig)
if err != nil {
log.Crit("Unable to create connection pool", "error", err)
os.Exit(1)
}
http.HandleFunc("/", urlHandler)
log.Info("Starting URL shortener on localhost:8080")
err = http.ListenAndServe("localhost:8080", nil)
if err != nil {
log.Crit("Unable to start web server", "error", err)
os.Exit(1)
}
}
func openPgxStdlib(config pgx.ConnPoolConfig) (*sql.DB, error) {
connPool, err := pgx.NewConnPool(config)
if err != nil {
return nil, err
}
return stdlib.OpenFromConnPool(connPool)
}
func createConnPool(t *testing.T, maxConnections int) *pgx.ConnPool {
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: maxConnections}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
return pool
}
// Get PostgreSQL pool
func (self *Database) GetPoll() (pool *pgx.ConnPool) {
pool, err := pgx.NewConnPool(self.getConfig())
if err != nil {
self.Log.Error(fmt.Sprintf("Unable to create connection pool to database: %v\n", err))
os.Exit(1)
}
self.Log.Info("starting Database service...")
return pool
}
func openTestClientMaxConns(t testing.TB, maxConnections int) *Client {
connPoolConfig := pgx.ConnPoolConfig{
ConnConfig: testConnConfig,
MaxConnections: maxConnections,
AfterConnect: PrepareStatements,
}
pool, err := pgx.NewConnPool(connPoolConfig)
if err != nil {
t.Fatal(err)
}
return NewClient(pool)
}
func initDatabase(dbHost string, dbUser string, dbPass string, dbName string, dbPort uint16, maxConnectionsInPool int) (*pgx.ConnPool, error) {
var successOrFailure string = "OK"
var config pgx.ConnPoolConfig
config.Host = dbHost
config.User = dbUser
config.Password = dbPass
config.Database = dbName
config.Port = dbPort
config.MaxConnections = maxConnectionsInPool
config.AfterConnect = func(conn *pgx.Conn) error {
worldSelectStmt = mustPrepare(conn, "worldSelectStmt", "SELECT id, randomNumber FROM World WHERE id = $1")
worldUpdateStmt = mustPrepare(conn, "worldUpdateStmt", "UPDATE World SET randomNumber = $1 WHERE id = $2")
fortuneSelectStmt = mustPrepare(conn, "fortuneSelectStmt", "SELECT id, message FROM Fortune")
// Disable synchronous commit for the current db connection
// as a performance optimization.
// See http://www.postgresql.org/docs/current/static/runtime-config-wal.html
// for details.
if _, err := conn.Exec("SET synchronous_commit TO OFF"); err != nil {
log.Fatalf("Error when disabling synchronous commit")
}
return nil
}
fmt.Println("--------------------------------------------------------------------------------------------")
connPool, err := pgx.NewConnPool(config)
if err != nil {
successOrFailure = "FAILED"
log.Println("Connecting to database ", dbName, " as user ", dbUser, " ", successOrFailure, ": \n ", err)
} else {
log.Println("Connecting to database ", dbName, " as user ", dbUser, ": ", successOrFailure)
log.Println("Fetching one record to test if db connection is valid...")
var w World
n := randomWorldNum()
if errPing := connPool.QueryRow("worldSelectStmt", n).Scan(&w.Id, &w.RandomNumber); errPing != nil {
log.Fatalf("Error scanning world row: %s", errPing)
}
log.Println("OK")
}
fmt.Println("--------------------------------------------------------------------------------------------")
return connPool, err
}
func TestNewConnPoolMaxConnectionsCannotBeLessThan2(t *testing.T) {
t.Parallel()
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 1}
pool, err := pgx.NewConnPool(config)
if err == nil {
pool.Close()
t.Fatal(`Expected NewConnPool to fail with "MaxConnections must be at least 2" error, but it succeeded`)
}
if err.Error() != "MaxConnections must be at least 2" {
t.Fatalf(`Expected NewConnPool to fail with "MaxConnections must be at least 2" error, but it failed with %v`, err)
}
}
// initial ConnPool of pgx
func (pgdb *PostgresDB) InitConnection() error {
//var pool *pgx.ConnPool
var err error
pgdb.Pool, err = pgx.NewConnPool(pgdb.poolConfig)
if err != nil {
log.Info("Unable to create connection pool", "error", err)
os.Exit(1)
}
log.Info("database connect sueecss")
return nil
}
func TestNewConnPoolDefaultsTo5MaxConnections(t *testing.T) {
t.Parallel()
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatal("Unable to establish connection pool")
}
defer pool.Close()
if n := pool.Stat().MaxConnections; n != 5 {
t.Fatalf("Expected pool to default to 5 max connections, but it was %d", n)
}
}
func setupTestDB(c *C, dbname string) *postgres.DB {
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
c.Fatal(err)
}
return postgres.New(pgxpool, nil)
}
func Open(conf *Conf, afterConn func(*pgx.Conn) error) (*DB, error) {
connConfig := pgx.ConnConfig{
Host: fmt.Sprintf("leader.%s.discoverd", conf.Service),
User: conf.User,
Database: conf.Database,
Password: conf.Password,
}
connPool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: connConfig,
AfterConnect: afterConn,
MaxConnections: 20,
AcquireTimeout: 30 * time.Second,
})
db := &DB{connPool, conf}
return db, err
}
func getPgxPool(dbURL string) (*pgx.ConnPool, error) {
pgxcfg, err := pgx.ParseURI(dbURL)
if err != nil {
return nil, err
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgxcfg,
AfterConnect: prepQue,
})
if err != nil {
return nil, err
}
return pgxpool, nil
}
func BenchmarkConnPool(b *testing.B) {
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 5}
pool, err := pgx.NewConnPool(config)
if err != nil {
b.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
b.ResetTimer()
for i := 0; i < b.N; i++ {
var conn *pgx.Conn
if conn, err = pool.Acquire(); err != nil {
b.Fatalf("Unable to acquire connection: %v", err)
}
pool.Release(conn)
}
}
func initDatabase(dbHost string, dbUser string, dbPass string, dbName string, dbPort uint16, maxConnectionsInPool int) (*pgx.ConnPool, error) {
var successOrFailure string = "OK"
var config pgx.ConnPoolConfig
config.Host = dbHost
config.User = dbUser
config.Password = dbPass
config.Database = dbName
config.Port = dbPort
config.MaxConnections = maxConnectionsInPool
config.AfterConnect = func(eachConn *pgx.Conn) error {
worldSelectStmt = mustPrepare(eachConn, "worldSelectStmt", "SELECT id, randomNumber FROM World WHERE id = $1")
worldUpdateStmt = mustPrepare(eachConn, "worldUpdateStmt", "UPDATE World SET randomNumber = $1 WHERE id = $2")
fortuneSelectStmt = mustPrepare(eachConn, "fortuneSelectStmt", "SELECT id, message FROM Fortune")
return nil
}
fmt.Println("--------------------------------------------------------------------------------------------")
connPool, err := pgx.NewConnPool(config)
if err != nil {
successOrFailure = "FAILED"
log.Println("Connecting to database ", dbName, " as user ", dbUser, " ", successOrFailure, ": \n ", err)
} else {
log.Println("Connecting to database ", dbName, " as user ", dbUser, ": ", successOrFailure)
log.Println("Fetching one record to test if db connection is valid...")
var w World
n := randomWorldNum()
if errPing := connPool.QueryRow("worldSelectStmt", n).Scan(&w.Id, &w.RandomNumber); errPing != nil {
log.Fatalf("Error scanning world row: %s", errPing)
}
log.Println("OK")
}
fmt.Println("--------------------------------------------------------------------------------------------")
return connPool, err
}
// NewQueClient create que client
func NewQueClient(dbURI string) (*que.Client, error) {
pgxcfg, err := pgx.ParseURI(dbURI)
if err != nil {
log.Fatal(err)
}
var qc *que.Client
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgxcfg,
AfterConnect: que.PrepareStatements,
})
if err != nil {
log.Fatal(err)
return qc, err
}
log.Println("create que client")
qc = que.NewClient(pgxpool)
return qc, nil
}
func createDB(t *testing.T, dbname string) *postgres.DB {
if dbname == "" {
dbname = "blobstoretest"
}
if err := pgtestutils.SetupPostgres(dbname); err != nil {
t.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
t.Fatal(err)
}
db := postgres.New(pgxpool, nil)
return db
}
func (s *S) SetUpSuite(c *C) {
dbname := "controllertest"
db := setupTestDB(c, dbname)
if err := migrateDB(db); err != nil {
c.Fatal(err)
}
// reconnect with que statements prepared now that schema is migrated
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "/var/run/postgresql",
Database: dbname,
},
AfterConnect: schema.PrepareStatements,
})
if err != nil {
c.Fatal(err)
}
db = postgres.New(pgxpool, nil)
ca, err := certgen.Generate(certgen.Params{IsCA: true})
if err != nil {
c.Fatal(err)
}
s.caCert = []byte(ca.PEM)
s.flac = newFakeLogAggregatorClient()
s.cc = tu.NewFakeCluster()
s.hc = handlerConfig{
db: db,
cc: s.cc,
lc: s.flac,
rc: newFakeRouter(),
keys: []string{authKey},
caCert: s.caCert,
}
handler := appHandler(s.hc)
s.srv = httptest.NewServer(handler)
client, err := controller.NewClient(s.srv.URL, authKey)
c.Assert(err, IsNil)
s.c = client
}
func initDatabase() error {
host, err := getCredentialString("benchmarkerPostgres", "host")
if err != nil {
return err
}
portFloat, err := getCredentialNumber("benchmarkerPostgres", "port")
if err != nil {
return err
}
port := uint16(portFloat)
dbName, err := getCredentialString("benchmarkerPostgres", "database")
if err != nil {
return err
}
user, err := getCredentialString("benchmarkerPostgres", "user")
if err != nil {
return err
}
password, err := getCredentialString("benchmarkerPostgres", "password")
if err != nil {
return err
}
connConfig := pgx.ConnConfig{
Host: host,
Port: port,
Database: dbName,
User: user,
Password: password,
}
poolConfig := pgx.ConnPoolConfig{
ConnConfig: connConfig,
MaxConnections: 3,
AfterConnect: prepareStatements,
}
connPool, err = pgx.NewConnPool(poolConfig)
if err != nil {
return err
}
return nil
}
func TestConnQueryLog(t *testing.T) {
logger := &testLogger{}
connConfig := pgx.ConnConfig{
Host: "127.0.0.1",
User: "******",
Password: "******",
Database: "pgx_test",
Logger: logger,
}
config := pgx.ConnPoolConfig{ConnConfig: connConfig}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
db, err := stdlib.OpenFromConnPool(pool)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer closeDB(t, db)
// clear logs from initial connection
logger.logs = []testLog{}
var n int64
err = db.QueryRow("select 1").Scan(&n)
if err != nil {
t.Fatalf("db.QueryRow unexpectedly failed: %v", err)
}
l := logger.logs[0]
if l.msg != "Query" {
t.Errorf("Expected to log Query, but got %v", l)
}
if !(l.ctx[0] == "sql" && l.ctx[1] == "select 1") {
t.Errorf("Expected to log Query with sql 'select 1', but got %v", l)
}
}
func NewPgxPool(url string) (*pgx.ConnPool, error) {
dbcfg := pgx.ConnPoolConfig{
MaxConnections: maxConnectionsFlag.Value(),
}
var err error
if url == "" {
dbcfg.ConnConfig, err = pgx.ParseEnvLibpq()
} else if strings.HasPrefix(url, "postgresql://") {
dbcfg.ConnConfig, err = pgx.ParseURI(url)
} else {
dbcfg.ConnConfig, err = pgx.ParseDSN(url)
}
if err != nil {
return nil, err
}
return pgx.NewConnPool(dbcfg)
}
