Load types using a single SQL query

When loading even a single type into pgx's type map, multiple SQL
queries are performed in series. Over a slow link, this is not ideal.
Worse, if multiple types are being registered, this is repeated multiple
times.

This commit changes the internal implementation of LoadType to use a
single SQL query. It also added a LoadTypes, which can retrieve type
mapping information for multiple types in a single SQL call.
Additionally, LoadTypes will recursively load any related types,
avoiding the need to explicitly list everything.

conn.go

@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -107,8 +108,10 @@ var (
 	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")
+var (
+	errDisabledStatementCache   = fmt.Errorf("cannot use QueryExecModeCacheStatement with disabled statement cache")
+	errDisabledDescriptionCache = fmt.Errorf("cannot use QueryExecModeCacheDescribe with disabled description cache")
+)
 
 // Connect establishes a connection with a PostgreSQL server with a connection string. See
 // pgconn.Connect for details.
@@ -843,7 +846,6 @@ func (c *Conn) getStatementDescription(
 	mode QueryExecMode,
 	sql string,
 ) (sd *pgconn.StatementDescription, err error) {
-
 	switch mode {
 	case QueryExecModeCacheStatement:
 		if c.statementCache == nil {
@@ -1393,3 +1395,216 @@ func (c *Conn) deallocateInvalidatedCachedStatements(ctx context.Context) error
 
 	return nil
 }
+
+/*
+buildLoadTypesSQL generates the correct query for retrieving type information.
+
+		pgVersion: the major version of the PostgreSQL server
+	    isCockroachDB: is this cockroachDB?
+		typeNames: the names of the types to load. If nil, load all types.
+*/
+func buildLoadTypesSQL(pgVersion int64, isCockroachDB bool, typeNames []string) string {
+	supportsMultirange := (pgVersion >= 14)
+	var typeNamesClause string
+	if typeNames == nil {
+		typeNamesClause = "IS NOT NULL"
+	} else {
+		typeNamesClause = "= ANY($1)"
+	}
+	parts := make([]string, 0, 10)
+
+	parts = append(parts, `
+WITH RECURSIVE
+selected_classes(oid,reltype) AS (
+    SELECT pg_class.oid, pg_class.reltype
+    FROM pg_catalog.pg_class
+        LEFT JOIN pg_catalog.pg_namespace n ON n.oid = pg_class.relnamespace
+    WHERE pg_catalog.pg_table_is_visible(pg_class.oid)
+      AND relname `, typeNamesClause, `
+UNION ALL
+    SELECT pg_class.oid, pg_class.reltype
+    FROM pg_class
+    INNER JOIN pg_namespace ON (pg_class.relnamespace = pg_namespace.oid)
+    WHERE nspname || '.' || relname `, typeNamesClause, `
+),
+selected_types(oid) AS (
+    SELECT reltype AS oid
+    FROM selected_classes
+UNION ALL
+    SELECT oid
+    FROM pg_type
+    WHERE typname `, typeNamesClause, `
+),
+relationships(parent, child, depth) AS (
+    SELECT DISTINCT 0::OID, selected_types.oid, 0
+    FROM selected_types
+UNION ALL
+    SELECT pg_type.oid AS parent, pg_attribute.atttypid AS child, 1
+    FROM selected_classes c
+    inner join pg_type ON (c.reltype = pg_type.oid)
+    inner join pg_attribute on (c.oid = pg_attribute.attrelid)
+UNION ALL
+    SELECT parent.oid, parent.typbasetype, relationships.depth + 1
+    FROM pg_type parent
+    INNER JOIN relationships ON (parent.oid = relationships.child)
+    WHERE parent.typtypmod = -1 AND parent.typbasetype != 0
+),
+composite AS (
+    SELECT pg_type.oid, ARRAY_AGG(attname ORDER BY attnum) AS attnames, ARRAY_AGG(atttypid ORDER BY ATTNUM) AS atttypids
+    FROM pg_attribute
+    INNER JOIN pg_class ON (pg_class.oid = pg_attribute.attrelid)
+    INNER JOIN pg_type ON (pg_type.oid = pg_class.reltype)
+    WHERE NOT attisdropped
+      AND attnum > 0
+    GROUP BY pg_type.oid
+)
+SELECT typname,
+       typtype,
+       typbasetype,
+       typelem,
+       pg_type.oid,`)
+	if supportsMultirange {
+		parts = append(parts, `
+       COALESCE(multirange.rngtypid, 0) AS rngtypid,`)
+	} else {
+		parts = append(parts, `
+       0 AS rngtypid,`)
+	}
+	parts = append(parts, `
+       COALESCE(pg_range.rngsubtype, 0) AS rngsubtype,
+       attnames, atttypids
+    FROM relationships
+    INNER JOIN pg_type ON (pg_type.oid IN ( relationships.child,relationships.parent) )
+    LEFT OUTER JOIN pg_range ON (pg_type.oid = pg_range.rngtypid)`)
+	if supportsMultirange {
+		parts = append(parts, `
+    LEFT OUTER JOIN pg_range multirange ON (pg_type.oid = multirange.rngmultitypid)`)
+	}
+
+	parts = append(parts, `
+    LEFT OUTER JOIN composite USING (oid)
+    WHERE typtype != 'b'`)
+	parts = append(parts, `
+    GROUP BY typname, typtype, typbasetype, typelem, pg_type.oid, pg_range.rngsubtype,`)
+	if supportsMultirange {
+		parts = append(parts, `
+        multirange.rngtypid,`)
+	}
+	parts = append(parts, `
+        attnames, atttypids
+    ORDER BY MAX(depth) desc, typname;`)
+	return strings.Join(parts, "")
+}
+
+// LoadTypes inspects the database for []typeNames and automatically registers all discovered
+// types. Any types referenced by these will also be included in the registration.
+func (c *Conn) LoadAndRegisterTypes(ctx context.Context, typeNames []string) error {
+	if typeNames == nil || len(typeNames) == 0 {
+		return fmt.Errorf("No type names were supplied.")
+	}
+	return c.loadAndRegisterTypes(ctx, typeNames, c.TypeMap())
+}
+
+func (c *Conn) loadAndRegisterTypes(ctx context.Context, typeNames []string, registerWith *pgtype.Map) error {
+	if registerWith == nil {
+		return fmt.Errorf("Type map must be supplied")
+	}
+	serverVersion, err := c.ServerVersion()
+	if err != nil {
+		return fmt.Errorf("Unexpected server version error: %w", err)
+	}
+	sql := buildLoadTypesSQL(serverVersion, c.IsCockroachDB(), typeNames)
+	var rows Rows
+	if typeNames == nil {
+		rows, err = c.Query(ctx, sql, QueryExecModeSimpleProtocol)
+	} else {
+		rows, err = c.Query(ctx, sql, QueryExecModeSimpleProtocol, typeNames)
+	}
+	if err != nil {
+		return fmt.Errorf("While generating load types query: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var oid uint32
+		var typeName, typtype string
+		var typbasetype, typelem uint32
+		var rngsubtype, rngtypid uint32
+		attnames := make([]string, 0, 0)
+		atttypids := make([]uint32, 0, 0)
+		err = rows.Scan(&typeName, &typtype, &typbasetype, &typelem, &oid, &rngtypid, &rngsubtype, &attnames, &atttypids)
+		if err != nil {
+			return fmt.Errorf("While scanning type information: %w", err)
+		}
+
+		switch typtype {
+		case "b": // array
+			dt, ok := c.TypeMap().TypeForOID(typelem)
+			if !ok {
+				return fmt.Errorf("array element OID not registered while loading %v", typeName)
+			}
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.ArrayCodec{ElementType: dt}})
+		case "c": // composite
+			var fields []pgtype.CompositeCodecField
+			for i, fieldName := range attnames {
+				//if fieldOID64, err = strconv.ParseUint(composite_fields[i+1], 10, 32); err != nil {
+				//	return nil, fmt.Errorf("While extracting OID used in composite field: %w", err)
+				//}
+				dt, ok := c.TypeMap().TypeForOID(atttypids[i])
+				if !ok {
+					return fmt.Errorf("unknown composite type field OID %v (%v)", atttypids[i], fieldName)
+				}
+				fields = append(fields, pgtype.CompositeCodecField{Name: fieldName, Type: dt})
+			}
+			if err != nil {
+				return fmt.Errorf("While parsing %v: %w", typeName, err)
+			}
+
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.CompositeCodec{Fields: fields}})
+		case "d": // domain
+			dt, ok := c.TypeMap().TypeForOID(typbasetype)
+			if !ok {
+				return fmt.Errorf("domain base type OID not registered for %v", typeName)
+			}
+
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: dt.Codec})
+		case "e": // enum
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.EnumCodec{}})
+		case "r": // range
+			dt, ok := c.TypeMap().TypeForOID(rngsubtype)
+			if !ok {
+				return fmt.Errorf("range element OID %v not registered for %v", rngsubtype, typeName)
+			}
+
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.RangeCodec{ElementType: dt}})
+		case "m": // multirange
+			dt, ok := c.TypeMap().TypeForOID(rngtypid)
+			if !ok {
+				return fmt.Errorf("multirange element OID %v not registered while loading %v", rngtypid, typeName)
+			}
+
+			registerWith.RegisterType(&pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.MultirangeCodec{ElementType: dt}})
+		default:
+			return fmt.Errorf("unknown typtype %v for %v", typtype, typeName)
+		}
+	}
+	return nil
+}
+
+func (conn *Conn) IsCockroachDB() bool {
+	return conn.PgConn().ParameterStatus("crdb_version") != ""
+}
+
+func (conn *Conn) ServerVersion() (int64, error) {
+	serverVersionStr := conn.PgConn().ParameterStatus("server_version")
+	serverVersionStr = regexp.MustCompile(`^[0-9]+`).FindString(serverVersionStr)
+	// if not PostgreSQL do nothing
+	if serverVersionStr == "" {
+		return 0, fmt.Errorf("Cannot identify server version in %q", serverVersionStr)
+	}
+
+	serverVersion, err := strconv.ParseInt(serverVersionStr, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("postgres version parsing failed: %w", err)
+	}
+	return serverVersion, nil
+}

pgtype/composite_test.go

@@ -10,11 +10,56 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestCompositeCodecTranscode(t *testing.T) {
+func TestCompositeCodecTranscodeWithLoadTypes(t *testing.T) {
 	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		_, err := conn.Exec(ctx, `drop domain if exists anotheruint64;
+drop type if exists ct_test;
+create domain anotheruint64 as numeric(20,0);
+
+create type ct_test as (
+  a text,
+  b int4,
+  c anotheruint64
+);`)
+		require.NoError(t, err)
+		defer conn.Exec(ctx, "drop type ct_test")
+		defer conn.Exec(ctx, "drop domain anotheruint64")
+
+		err = conn.LoadAndRegisterTypes(ctx, []string{"ct_test"})
+		require.NoError(t, err)
+
+		formats := []struct {
+			name string
+			code int16
+		}{
+			{name: "TextFormat", code: pgx.TextFormatCode},
+			{name: "BinaryFormat", code: pgx.BinaryFormatCode},
+		}
+
+		for _, format := range formats {
+			var a string
+			var b int32
+			var c uint64
 
+			err := conn.QueryRow(ctx, "select $1::ct_test", pgx.QueryResultFormats{format.code},
+				pgtype.CompositeFields{"hi", int32(42), uint64(123)},
+			).Scan(
+				pgtype.CompositeFields{&a, &b, &c},
+			)
+			require.NoErrorf(t, err, "%v", format.name)
+			require.EqualValuesf(t, "hi", a, "%v", format.name)
+			require.EqualValuesf(t, 42, b, "%v", format.name)
+			require.EqualValuesf(t, 123, c, "%v", format.name)
+		}
+	})
+}
+
+func TestCompositeCodecTranscode(t *testing.T) {
+	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
+
+	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
 		_, err := conn.Exec(ctx, `drop type if exists ct_test;
 
 create type ct_test as (
@@ -94,7 +139,6 @@ func TestCompositeCodecTranscodeStruct(t *testing.T) {
 	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
-
 		_, err := conn.Exec(ctx, `drop type if exists point3d;
 
 create type point3d as (
@@ -131,7 +175,6 @@ func TestCompositeCodecTranscodeStructWrapper(t *testing.T) {
 	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
-
 		_, err := conn.Exec(ctx, `drop type if exists point3d;
 
 create type point3d as (
@@ -172,7 +215,6 @@ func TestCompositeCodecDecodeValue(t *testing.T) {
 	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
-
 		_, err := conn.Exec(ctx, `drop type if exists point3d;
 
 create type point3d as (
@@ -217,7 +259,6 @@ func TestCompositeCodecTranscodeStructWrapperForTable(t *testing.T) {
 	skipCockroachDB(t, "Server does not support composite types (see https://github.com/cockroachdb/cockroach/issues/27792)")
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
-
 		_, err := conn.Exec(ctx, `drop table if exists point3d;
 
 create table point3d (