Repository: jackc/pgx

Issue:
Convert pgxpool conn to stdlib connection
I was wondering if it's possible to convert a pgxpool Conn into a `*sql.DB` to pass to other libraries that don't support pgx or pgxpool?

Discussion:
@jackc: No. But it can work the other way. You can use https://pkg.go.dev/database/sql#Conn.Raw to go from `*sql.Conn` to `*pgx.Conn`.@aranw: Ah I was really hoping to go from `pgxpool` to `sql.DB` as I want to use `pgxpool` as my main connection pool management and then pass `sql.Conn/DB` into libraries that don't support pgx properly. I've noticed odd behaviour with the default `database/sql` pooling @dio: @aranw, I kinda think the same since I have a utility to be used against MySQL and Postgres (hence prefer to pass `*sql.DB`). Sorry if this is unrelated, but did you find `pgxpool` does a better job? I'd really keen to know more. Thanks!@vec715: this issue is still relevant for me, the solution has not been found yet?@owenthereal: Any progress on this? I want to pass around the pgxpool object for queries but some libraries only work with *sql.DB@owenthereal: One solution is that we allow creating a *sql.DB from a pgxpool in stdlib@RohanNg: This commit https://github.com/jackc/pgx/commit/ffae1b134596a9a111fdb89b1bd09febad7e2dae removes the ability to convert from pgxpool to sql.DB. Do we have any context into why it was removed?@jackc: >This commit ffae1b1 removes the ability to convert from pgxpool to sql.DB. Do we have any context into why it was removed?

It worked, but there were a lot of confusing edges with having two pools managing one set of connections. e.g. If both pools had the same number max connections then if the `*sql.DB` acquired all of them then the pgxpool would be starved of connections.

If the pgxpool to `*sql.DB` feature was to be resurrected it would probably make sense to limit the `*sql.DB` to a single connection. The API would be more like the pool's `Acquire()` / `Release()` only it would acquire and release a `*sql.DB` instead of a `*pgxpool.Conn`. i.e. It would explicitly be designed for a short lifetime.@yufan-xue: > No. But it can work the other way. You can use https://pkg.go.dev/database/sql#Conn.Raw to go from `*sql.Conn` to `*pgx.Conn`.

I tried the method, but it doesn't work and failed silently...
```
err := conn.Raw(func(driverConn interface{}) error {
  pgxConn, ok := driverConn.(*pgx.Conn)
  if !ok {
    return fmt.Errorf("super bad, expected a *pgx.Conn, but have %T", driverConn)
  }
  do thing with pgxConn   
})
``` 

@jackc @jackc: From the package docs -- https://pkg.go.dev/github.com/jackc/pgx/v4/stdlib.

```go
// Given db is a *sql.DB
conn, err := db.Conn(context.Background())
if err != nil {
	// handle error from acquiring connection from DB pool
}

err = conn.Raw(func(driverConn interface{}) error {
	conn := driverConn.(*stdlib.Conn).Conn() // conn is a *pgx.Conn
	// Do pgx specific stuff with conn
	conn.CopyFrom(...)
	return nil
})
if err != nil {
	// handle error that occurred while using *pgx.Conn
}
```

@StevenACoffman: If it helps, our use case is different from what was assumed.

When our app starts up, we inject a pgxpool, which has a specific set of max database connections configured. 

At app startup, we detect if the schema migrations have been performed and autoupdate them (typically on a local developer machine, not in prod). However, the schema migrations tool (golang-migrate) requires a `*sql.DB` and we are stuck with a `pgx/v5/pgxpool` that cannot be converted.

Workaround for now is to reorder things so we first create `*sql.DB`, auto-migrate, close `*sql.DB`, and then create  `pgx/v5/pgxpool`.  That way we don't have a `*sql.DB` and a `*pgx/v5/pgxpool.Pool` at the same time eating double the max connections.@bodokaiser: This feature is really important!

We need to implement tracing (via pgx) to debug connection issues but cannot use the pgx interface because entgo requires a `*sql.DB`.@redbaron: Exposing pgxpoool as `*sql.DB` could help us as well. Thinking about it, it can be done reliably if `database/sql` doesn't try to maintain internal pool, that is max idle connections is set to 0  , max open cons to infinity and ignoring all other options (max lifetime, max idle time).  It should probably be implemented by a small `pgxpool://` driver which wraps pool and translates [Conn.Close](https://pkg.go.dev/database/sql#Conn.Close) to [Release](https://pkg.go.dev/github.com/jackc/pgx/v5/pgxpool#Conn.Release)@jackc: >  Thinking about it, it can be done reliably if database/sql doesn't try to maintain internal pool, that is max idle connections is set to 0 , max open cons to infinity and ignoring all other options (max lifetime, max idle time).

That might be workable. I think the previous implementation of this feature predated database/sql exposing those knobs.@kcgopal: I haven't gone through the internal working fo pgxpool.  As per my understanding "Acquire" allots one entry from pool and marks it as unavailable and "Release" gives the resource back to the pool.  If that is the case Isn't it possible to have couple more function in pgxpool which is like Aquire & Release, but returns and releases *sql.DB from pool.  That way your concern about too pool managing one set of connections won't be there as the pgxpool will be in control of allocating  for both.  @lzakharov: > Thinking about it, it can be done reliably if database/sql doesn't try to maintain internal pool, that is max idle connections is set to 0 , max open cons to infinity and ignoring all other options (max lifetime, max idle time). 

@redbaron literally now I'm implementing a similar scheme and so far it looks like a workable option. Just changed [this driver.Connector implementation](https://github.com/jackc/pgx/blob/master/stdlib/sql.go#L199) to work with `*pgxpool.Pool`. The only thing I don't like is the inability to return the correct `DBStats` from the pool `*sql.DB`, but it doesn't look like a problem, at least in my case.@redbaron: @lzakharov , sounds great! do you think it will be mergeable as something like `pgx/v5/pgxpool/stdlib`?@lzakharov: @redbaron yeap, or it can also be added directly to the `pgx/v5/stdlib` with constructor like `OpenDBFromPool`. I'll try to provide this implementation soon.@zcubbs: this feature is a must!@gluck1986: guys, i'v found this 
```go
sql.OpenDB(stdlib.GetConnector(*pool.Config().ConnConfig))
```@SimonWoidig: > guys, i'v found this
> 
> sql.OpenDB(stdlib.GetConnector(*pool.Config().ConnConfig))

Even better would be to use:
```go
sql.OpenDB(stdlib.GetPoolConnector(pool))
```

Solution patch:
diff --git a/conn.go b/conn.go
index 7c7081b4..b867acfe 100644
--- a/conn.go
+++ b/conn.go
@@ -99,8 +99,12 @@ func (ident Identifier) Sanitize() string {
 	return strings.Join(parts, ".")
 }
 
-// ErrNoRows occurs when rows are expected but none are returned.
-var ErrNoRows = errors.New("no rows in result set")
+var (
+	// ErrNoRows occurs when rows are expected but none are returned.
+	ErrNoRows = errors.New("no rows in result set")
+	// ErrTooManyRows occurs when more rows than expected are returned.
+	ErrTooManyRows = errors.New("too many rows in result set")
+)
 
 var errDisabledStatementCache = fmt.Errorf("cannot use QueryExecModeCacheStatement with disabled statement cache")
 var errDisabledDescriptionCache = fmt.Errorf("cannot use QueryExecModeCacheDescribe with disabled description cache")
diff --git a/rows.go b/rows.go
index ab8a12f8..c4c7ee0f 100644
--- a/rows.go
+++ b/rows.go
@@ -48,7 +48,7 @@ type Rows interface {
 	// Callers should check rows.Err() after rows.Next() returns false to detect
 	// whether result-set reading ended prematurely due to an error. See
 	// Conn.Query for details.
-        //
+	//
 	// For simpler error handling, consider using the higher-level pgx v5
 	// CollectRows() and ForEachRow() helpers instead.
 	Next() bool
@@ -465,6 +465,39 @@ func CollectOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) {
 	return value, rows.Err()
 }
 
+// CollectExactlyOneRow calls fn for the first row in rows and returns the result.
+//   - If no rows are found returns an error where errors.Is(ErrNoRows) is true.
+//   - If more than 1 row is found returns an error where errors.Is(ErrTooManyRows) is true.
+func CollectExactlyOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) {
+	defer rows.Close()
+
+	var (
+		err   error
+		value T
+	)
+
+	if !rows.Next() {
+		if err = rows.Err(); err != nil {
+			return value, err
+		}
+
+		return value, ErrNoRows
+	}
+
+	value, err = fn(rows)
+	if err != nil {
+		return value, err
+	}
+
+	if rows.Next() {
+		var zero T
+
+		return zero, ErrTooManyRows
+	}
+
+	return value, rows.Err()
+}
+
 // RowTo returns a T scanned from row.
 func RowTo[T any](row CollectableRow) (T, error) {
 	var value T
diff --git a/stdlib/sql.go b/stdlib/sql.go
index 97ecc9b2..6a081e21 100644
--- a/stdlib/sql.go
+++ b/stdlib/sql.go
@@ -14,6 +14,18 @@
 //	  return err
 //	}
 //
+// Or from a *pgxpool.Pool.
+//
+//	pool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
+//	if err != nil {
+//	  return err
+//	}
+//
+//	db, err := stdlib.OpenDBFromPool(pool)
+//	if err != nil {
+//	  return err
+//	}
+//
 // Or a pgx.ConnConfig can be used to set configuration not accessible via connection string. In this case the
 // pgx.ConnConfig must first be registered with the driver. This registration returns a connection string which is used
 // with sql.Open.
@@ -74,6 +86,7 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 // Only intrinsic types should be binary format with database/sql.
@@ -125,14 +138,14 @@ func contains(list []string, y string) bool {
 type OptionOpenDB func(*connector)
 
 // OptionBeforeConnect provides a callback for before connect. It is passed a shallow copy of the ConnConfig that will
-// be used to connect, so only its immediate members should be modified.
+// be used to connect, so only its immediate members should be modified. Used only if db is opened with *pgx.ConnConfig.
 func OptionBeforeConnect(bc func(context.Context, *pgx.ConnConfig) error) OptionOpenDB {
 	return func(dc *connector) {
 		dc.BeforeConnect = bc
 	}
 }
 
-// OptionAfterConnect provides a callback for after connect.
+// OptionAfterConnect provides a callback for after connect. Used only if db is opened with *pgx.ConnConfig.
 func OptionAfterConnect(ac func(context.Context, *pgx.Conn) error) OptionOpenDB {
 	return func(dc *connector) {
 		dc.AfterConnect = ac
@@ -191,13 +204,38 @@ func GetConnector(config pgx.ConnConfig, opts ...OptionOpenDB) driver.Connector
 	return c
 }
 
+func GetPoolConnector(pool *pgxpool.Pool, opts ...OptionOpenDB) driver.Connector {
+	c := connector{
+		pool:         pool,
+		ResetSession: func(context.Context, *pgx.Conn) error { return nil }, // noop reset session by default
+		driver:       pgxDriver,
+	}
+
+	for _, opt := range opts {
+		opt(&c)
+	}
+
+	return c
+}
+
 func OpenDB(config pgx.ConnConfig, opts ...OptionOpenDB) *sql.DB {
 	c := GetConnector(config, opts...)
 	return sql.OpenDB(c)
 }
 
+// OpenDBFromPool creates a new *sql.DB from the given *pgxpool.Pool. Note that this method automatically sets the
+// maximum number of idle connections in *sql.DB to zero, since they must be managed from the *pgxpool.Pool. This is
+// required to avoid acquiring all the connections from the pgxpool and starving any direct users of the pgxpool.
+func OpenDBFromPool(pool *pgxpool.Pool, opts ...OptionOpenDB) *sql.DB {
+	c := GetPoolConnector(pool, opts...)
+	db := sql.OpenDB(c)
+	db.SetMaxIdleConns(0)
+	return db
+}
+
 type connector struct {
 	pgx.ConnConfig
+	pool          *pgxpool.Pool
 	BeforeConnect func(context.Context, *pgx.ConnConfig) error // function to call before creation of every new connection
 	AfterConnect  func(context.Context, *pgx.Conn) error       // function to call after creation of every new connection
 	ResetSession  func(context.Context, *pgx.Conn) error       // function is called before a connection is reused
@@ -207,25 +245,52 @@ type connector struct {
 // Connect implement driver.Connector interface
 func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
 	var (
-		err  error
-		conn *pgx.Conn
+		connConfig pgx.ConnConfig
+		conn       *pgx.Conn
+		close      func(context.Context) error
+		err        error
 	)
 
-	// Create a shallow copy of the config, so that BeforeConnect can safely modify it
-	connConfig := c.ConnConfig
-	if err = c.BeforeConnect(ctx, &connConfig); err != nil {
-		return nil, err
-	}
+	if c.pool == nil {
+		// Create a shallow copy of the config, so that BeforeConnect can safely modify it
+		connConfig = c.ConnConfig
 
-	if conn, err = pgx.ConnectConfig(ctx, &connConfig); err != nil {
-		return nil, err
-	}
+		if err = c.BeforeConnect(ctx, &connConfig); err != nil {
+			return nil, err
+		}
 
-	if err = c.AfterConnect(ctx, conn); err != nil {
-		return nil, err
+		if conn, err = pgx.ConnectConfig(ctx, &connConfig); err != nil {
+			return nil, err
+		}
+
+		if err = c.AfterConnect(ctx, conn); err != nil {
+			return nil, err
+		}
+
+		close = conn.Close
+	} else {
+		var pconn *pgxpool.Conn
+
+		pconn, err = c.pool.Acquire(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		conn = pconn.Conn()
+
+		close = func(_ context.Context) error {
+			pconn.Release()
+			return nil
+		}
 	}
 
-	return &Conn{conn: conn, driver: c.driver, connConfig: connConfig, resetSessionFunc: c.ResetSession}, nil
+	return &Conn{
+		conn:             conn,
+		close:            close,
+		driver:           c.driver,
+		connConfig:       connConfig,
+		resetSessionFunc: c.ResetSession,
+	}, nil
 }
 
 // Driver implement driver.Connector interface
@@ -302,6 +367,7 @@ func (dc *driverConnector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	c := &Conn{
 		conn:             conn,
+		close:            conn.Close,
 		driver:           dc.driver,
 		connConfig:       *connConfig,
 		resetSessionFunc: func(context.Context, *pgx.Conn) error { return nil },
@@ -326,6 +392,7 @@ func UnregisterConnConfig(connStr string) {
 
 type Conn struct {
 	conn                 *pgx.Conn
+	close                func(context.Context) error
 	psCount              int64 // Counter used for creating unique prepared statement names
 	driver               *Driver
 	connConfig           pgx.ConnConfig
@@ -361,7 +428,7 @@ func (c *Conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, e
 func (c *Conn) Close() error {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
-	return c.conn.Close(ctx)
+	return c.close(ctx)
 }
 
 func (c *Conn) Begin() (driver.Tx, error) {



Test patch:
diff --git a/rows_test.go b/rows_test.go
index b2d1137a..be221ea4 100644
--- a/rows_test.go
+++ b/rows_test.go
@@ -274,6 +274,45 @@ func TestCollectOneRowPrefersPostgreSQLErrorOverErrNoRows(t *testing.T) {
 	})
 }
 
+func TestCollectExactlyOneRow(t *testing.T) {
+	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		rows, _ := conn.Query(ctx, `select 42`)
+		n, err := pgx.CollectExactlyOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
+			var n int32
+			err := row.Scan(&n)
+			return n, err
+		})
+		assert.NoError(t, err)
+		assert.Equal(t, int32(42), n)
+	})
+}
+
+func TestCollectExactlyOneRowNotFound(t *testing.T) {
+	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		rows, _ := conn.Query(ctx, `select 42 where false`)
+		n, err := pgx.CollectExactlyOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
+			var n int32
+			err := row.Scan(&n)
+			return n, err
+		})
+		assert.ErrorIs(t, err, pgx.ErrNoRows)
+		assert.Equal(t, int32(0), n)
+	})
+}
+
+func TestCollectExactlyOneRowExtraRows(t *testing.T) {
+	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		rows, _ := conn.Query(ctx, `select n from generate_series(42, 99) n`)
+		n, err := pgx.CollectExactlyOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
+			var n int32
+			err := row.Scan(&n)
+			return n, err
+		})
+		assert.ErrorIs(t, err, pgx.ErrTooManyRows)
+		assert.Equal(t, int32(0), n)
+	})
+}
+
 func TestRowTo(t *testing.T) {
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
 		rows, _ := conn.Query(ctx, `select n from generate_series(0, 99) n`)
diff --git a/stdlib/sql_test.go b/stdlib/sql_test.go
index aa8a613b..93265726 100644
--- a/stdlib/sql_test.go
+++ b/stdlib/sql_test.go
@@ -18,6 +18,7 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jackc/pgx/v5/stdlib"
 	"github.com/jackc/pgx/v5/tracelog"
 	"github.com/stretchr/testify/assert"
@@ -166,6 +167,17 @@ func TestSQLOpen(t *testing.T) {
 	}
 }
 
+func TestSQLOpenFromPool(t *testing.T) {
+	pool, err := pgxpool.New(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	db := stdlib.OpenDBFromPool(pool)
+	ensureDBValid(t, db)
+
+	db.Close()
+}
+
 func TestNormalLifeCycle(t *testing.T) {
 	db := openDB(t)
 	defer closeDB(t, db)