// Do a simple query to ensure the connection is still usable
func ensureConnValid(t *testing.T, conn *pgx.Conn) {
var sum, rowCount int32
rows, err := conn.Query("select generate_series(1,$1)", 10)
if err != nil {
t.Fatalf("conn.Query failed: ", err)
}
defer rows.Close()
for rows.Next() {
var n int32
rows.Scan(&n)
sum += n
rowCount++
}
if rows.Err() != nil {
t.Fatalf("conn.Query failed: ", err)
}
if rowCount != 10 {
t.Error("Select called onDataRow wrong number of times")
}
if sum != 55 {
t.Error("Wrong values returned")
}
}
func mustExec(t testing.TB, conn *pgx.Conn, sql string, arguments ...interface{}) (commandTag pgx.CommandTag) {
var err error
if commandTag, err = conn.Exec(sql, arguments...); err != nil {
t.Fatalf("Exec unexpectedly failed with %v: %v", sql, err)
}
return
}
func mustPrepare(t testing.TB, conn *pgx.Conn, name, sql string) *pgx.PreparedStatement {
ps, err := conn.Prepare(name, sql)
if err != nil {
t.Fatalf("Could not prepare %v: %v", name, err)
}
return ps
}
// afterConnect creates the prepared statements that this application uses
func afterConnect(conn *pgx.Conn) (err error) {
_, err = conn.Prepare("getUrl", `
select url from shortened_urls where id=$1
`)
if err != nil {
return
}
_, err = conn.Prepare("deleteUrl", `
delete from shortened_urls where id=$1
`)
if err != nil {
return
}
// There technically is a small race condition in doing an upsert with a CTE
// where one of two simultaneous requests to the shortened URL would fail
// with a unique index violation. As the point of this demo is pgx usage and
// not how to perfectly upsert in PostgreSQL it is deemed acceptable.
_, err = conn.Prepare("putUrl", `
with upsert as (
update shortened_urls
set url=$2
where id=$1
returning *
)
insert into shortened_urls(id, url)
select $1, $2 where not exists(select 1 from upsert)
`)
return
}
func TestPoolReleaseDiscardsDeadConnections(t *testing.T) {
t.Parallel()
maxConnections := 3
pool := createConnPool(t, maxConnections)
defer pool.Close()
var c1, c2 *pgx.Conn
var err error
var stat pgx.ConnPoolStat
if c1, err = pool.Acquire(); err != nil {
t.Fatalf("Unexpected error acquiring connection: %v", err)
}
defer func() {
if c1 != nil {
pool.Release(c1)
}
}()
if c2, err = pool.Acquire(); err != nil {
t.Fatalf("Unexpected error acquiring connection: %v", err)
}
defer func() {
if c2 != nil {
pool.Release(c2)
}
}()
if _, err = c2.Exec("select pg_terminate_backend($1)", c1.Pid); err != nil {
t.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
}
// do something with the connection so it knows it's dead
rows, _ := c1.Query("select 1")
rows.Close()
if rows.Err() == nil {
t.Fatal("Expected error but none occurred")
}
if c1.IsAlive() {
t.Fatal("Expected connection to be dead but it wasn't")
}
stat = pool.Stat()
if stat.CurrentConnections != 2 {
t.Fatalf("Unexpected CurrentConnections: %v", stat.CurrentConnections)
}
if stat.AvailableConnections != 0 {
t.Fatalf("Unexpected AvailableConnections: %v", stat.CurrentConnections)
}
pool.Release(c1)
c1 = nil // so it doesn't get released again by the defer
stat = pool.Stat()
if stat.CurrentConnections != 1 {
t.Fatalf("Unexpected CurrentConnections: %v", stat.CurrentConnections)
}
if stat.AvailableConnections != 0 {
t.Fatalf("Unexpected AvailableConnections: %v", stat.CurrentConnections)
}
}
func closeConn(t testing.TB, conn *pgx.Conn) {
err := conn.Close()
if err != nil {
t.Fatalf("conn.Close unexpectedly failed: %v", err)
}
}