From 37c3f157bcbdbd30b17885deee87efd1d1baaa53 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Thu, 14 Apr 2022 11:50:12 -0500
Subject: [PATCH] Add Hijack from v5

---
 pgxpool/conn.go      | 16 ++++++++++++++++
 pgxpool/pool_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/pgxpool/conn.go b/pgxpool/conn.go
index 0b59d741..036c728a 100644
--- a/pgxpool/conn.go
+++ b/pgxpool/conn.go
@@ -46,6 +46,22 @@ func (c *Conn) Release() {
 	}()
 }
 
+// Hijack assumes ownership of the connection from the pool. Caller is responsible for closing the connection. Hijack
+// will panic if called on an already released or hijacked connection.
+func (c *Conn) Hijack() *pgx.Conn {
+	if c.res == nil {
+		panic("cannot hijack already released or hijacked connection")
+	}
+
+	conn := c.Conn()
+	res := c.res
+	c.res = nil
+
+	res.Hijack()
+
+	return conn
+}
+
 func (c *Conn) Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error) {
 	return c.Conn().Exec(ctx, sql, arguments...)
 }
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index 072ca053..42e029e1 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -115,6 +115,32 @@ func TestPoolAcquireAndConnRelease(t *testing.T) {
 	c.Release()
 }
 
+func TestPoolAcquireAndConnHijack(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	pool, err := pgxpool.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer pool.Close()
+
+	c, err := pool.Acquire(ctx)
+	require.NoError(t, err)
+
+	connsBeforeHijack := pool.Stat().TotalConns()
+
+	conn := c.Hijack()
+	defer conn.Close(ctx)
+
+	connsAfterHijack := pool.Stat().TotalConns()
+	require.Equal(t, connsBeforeHijack-1, connsAfterHijack)
+
+	var n int32
+	err = conn.QueryRow(ctx, `select 1`).Scan(&n)
+	require.NoError(t, err)
+	require.Equal(t, int32(1), n)
+}
+
 func TestPoolAcquireFunc(t *testing.T) {
 	t.Parallel()