From 089aee39430e5d163c13c3c0438e220387d661d2 Mon Sep 17 00:00:00 2001
From: Nick Farrell <nick.farrell@aiven.io>
Date: Mon, 1 Jul 2024 08:12:04 +1000
Subject: [PATCH] Allow pgxpool to more easily use LoadTypes

While LoadTypes is powerful on its own, retriving all salient type
information in a single query, it is particularly powerful and useful
when users of pgxpool can use this.

To streamline its use, provide a helper function in pgxpool suitable for
plugging directly into the AfterConnect configuration settings. It will
load the types and register them, also reusing them if appropriate.
---
 pgxpool/pool.go | 68 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 63 insertions(+), 5 deletions(-)

diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index fdcba7241..61ff0573a 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"os"
 	"runtime"
 	"strconv"
 	"sync"
@@ -12,14 +13,17 @@ import (
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/puddle/v2"
 )
 
-var defaultMaxConns = int32(4)
-var defaultMinConns = int32(0)
-var defaultMaxConnLifetime = time.Hour
-var defaultMaxConnIdleTime = time.Minute * 30
-var defaultHealthCheckPeriod = time.Minute
+var (
+	defaultMaxConns          = int32(4)
+	defaultMinConns          = int32(0)
+	defaultMaxConnLifetime   = time.Hour
+	defaultMaxConnIdleTime   = time.Minute * 30
+	defaultHealthCheckPeriod = time.Minute
+)
 
 type connResource struct {
 	conn       *pgx.Conn
@@ -102,6 +106,60 @@ type Pool struct {
 	closeChan chan struct{}
 }
 
+// LoadTypesAfterConnect is suitable for assigning to the AfterConnect configuration setting.
+// It will automatically load the named types for each connection in an efficient manner,
+// performing a single query to the database backend. The underlying call to pgx.LoadTypes
+// is smart enough to also retrieve any related types required to support the definition of the
+// named types.
+// If reuseTypeMap is enabled, it is assumed that the OID mapping is stable across all database
+// backends in this pool, resulting in only needing to query when creating the initial connection;
+// subsequent connections will reuse the same OID type mapping.
+// Because it is not always possible for a client to know the database topology in the final usage
+// context, PGXPOOL_REUSE_TYPEMAP, when given a value of y or n, will take precedence over this argument.
+func LoadTypesAfterConnect(typeNames []string, reuseTypeMap bool) func(context.Context, *pgx.Conn) error {
+	switch os.Getenv("PGXPOOL_REUSE_TYPEMAP") {
+	case "y":
+		reuseTypeMap = true
+	case "n":
+		reuseTypeMap = false
+	}
+	if reuseTypeMap {
+		mutex := new(sync.Mutex)
+		var types []*pgtype.Type
+		return func(ctx context.Context, conn *pgx.Conn) error {
+			if types != nil {
+				// avoid acquiring the mutex if the types are already available
+				conn.TypeMap().RegisterTypes(types)
+				return nil
+			}
+			mutex.Lock()
+			defer mutex.Unlock()
+			var err error
+
+			// types may have become available while waiting for the mutex
+			if types != nil {
+				conn.TypeMap().RegisterTypes(types)
+				return nil
+			}
+			types, err = conn.LoadTypes(ctx, typeNames)
+			if err != nil {
+				types = nil
+				return err
+			}
+			conn.TypeMap().RegisterTypes(types)
+			return nil
+		}
+	}
+	return func(ctx context.Context, conn *pgx.Conn) error {
+		types, err := conn.LoadTypes(ctx, typeNames)
+		if err != nil {
+			return err
+		}
+		conn.TypeMap().RegisterTypes(types)
+		return nil
+	}
+}
+
 // Config is the configuration struct for creating a pool. It must be created by [ParseConfig] and then it can be
 // modified.
 type Config struct {