From 8ae83b19f7d6a2a27ac3b1a2664dc3c61a90cf46 Mon Sep 17 00:00:00 2001
From: Maxim Ivanov <hi+github@yamlcoder.me>
Date: Sun, 12 Apr 2020 22:33:33 +0100
Subject: [PATCH] Add EncodeRow helpers

Also extend example to show how EncodeRow can be used
to create binary encoders for composite type
---
 composite_test.go | 47 +++++++++++++++++++++++++++++++++--------------
 convert.go        | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 14 deletions(-)

diff --git a/composite_test.go b/composite_test.go
index ffa7d479..d0c48f6e 100644
--- a/composite_test.go
+++ b/composite_test.go
@@ -41,11 +41,32 @@ func (dst *MyType) DecodeBinary(ci *pgtype.ConnInfo, src []byte) error {
 	return nil
 }
 
-func Example_compositeTypes() {
-	conn, err := pgx.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
+func (src MyType) EncodeBinary(ci *pgtype.ConnInfo, buf []byte) (newBuf []byte, err error) {
+	a := pgtype.Int4{src.a, pgtype.Present}
+	var b pgtype.Text
+	if src.b != nil {
+		b = pgtype.Text{*src.b, pgtype.Present}
+	} else {
+		b = pgtype.Text{Status: pgtype.Null}
+	}
+
+	return pgtype.EncodeRow(ci, buf, &a, &b)
+}
+
+func ptrS(s string) *string {
+	return &s
+}
+
+func E(err error) {
 	if err != nil {
 		panic(err)
 	}
+}
+
+func Example_compositeTypes() {
+	conn, err := pgx.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
+	E(err)
+
 	defer conn.Close(context.Background())
 	_, err = conn.Exec(context.Background(), `drop type if exists mytype;
 
@@ -53,22 +74,21 @@ create type mytype as (
   a int4,
   b text
 );`)
-	if err != nil {
-		panic(err)
-	}
+	E(err)
 	defer conn.Exec(context.Background(), "drop type mytype")
 
 	var result *MyType
-	if err = conn.QueryRow(context.Background(), "select (1,'foo')::mytype", pgx.QueryResultFormats{pgx.BinaryFormatCode}).Scan(&result); err != nil {
-		panic(err)
-	}
+
+	// Demonstrates both passing and reading back composite values
+	err = conn.QueryRow(context.Background(), "select $1::mytype",
+		pgx.QueryResultFormats{pgx.BinaryFormatCode}, MyType{1, ptrS("foo")}).Scan(&result)
+	E(err)
 
 	fmt.Printf("First row: a=%d b=%s\n", result.a, *result.b)
 
 	// Because we scan into &*MyType, NULLs are handled generically by assigning nil to result
-	if err = conn.QueryRow(context.Background(), "select NULL::mytype", pgx.QueryResultFormats{pgx.BinaryFormatCode}).Scan(&result); err != nil {
-		panic(err)
-	}
+	err = conn.QueryRow(context.Background(), "select NULL::mytype", pgx.QueryResultFormats{pgx.BinaryFormatCode}).Scan(&result)
+	E(err)
 
 	fmt.Printf("Second row: %v\n", result)
 
@@ -77,9 +97,8 @@ create type mytype as (
 	var a int
 	var b *string
 
-	if err = conn.QueryRow(context.Background(), "select (2, 'bar')::mytype", pgx.QueryResultFormats{pgx.BinaryFormatCode}).Scan(pgtype.ROW(&isNull, &a, &b)); err != nil {
-		panic(err)
-	}
+	err = conn.QueryRow(context.Background(), "select (2, 'bar')::mytype", pgx.QueryResultFormats{pgx.BinaryFormatCode}).Scan(pgtype.ROW(&isNull, &a, &b))
+	E(err)
 
 	fmt.Printf("Adhoc: isNull=%v a=%d b=%s", isNull, a, *b)
 
diff --git a/convert.go b/convert.go
index 8157358b..d22a714f 100644
--- a/convert.go
+++ b/convert.go
@@ -5,6 +5,7 @@ import (
 	"reflect"
 	"time"
 
+	"github.com/jackc/pgio"
 	errors "golang.org/x/xerrors"
 )
 
@@ -471,6 +472,38 @@ func ScanRowValue(ci *ConnInfo, src []byte, dst ...Value) error {
 	return nil
 }
 
+// EncodeRow builds a binary representation of row values (row(), composite types)
+func EncodeRow(ci *ConnInfo, buf []byte, fields ...Value) (newBuf []byte, err error) {
+	fieldBytes := make([]byte, 0, 128)
+
+	newBuf = pgio.AppendUint32(buf, uint32(len(fields)))
+	for _, f := range fields {
+		dt, ok := ci.DataTypeForValue(f)
+		if !ok {
+			return nil, errors.Errorf("Unknown OID for %s", f)
+		}
+		newBuf = pgio.AppendUint32(newBuf, dt.OID)
+
+		if f.Get() != nil {
+			binaryEncoder, ok := f.(BinaryEncoder)
+			if !ok {
+				return nil, errors.Errorf("record field doesn't implement binary encoding: %s", reflect.TypeOf(f).Name())
+			}
+			fieldBytes, err = binaryEncoder.EncodeBinary(ci, fieldBytes[:0])
+			if err != nil {
+				return nil, err
+			}
+
+			newBuf = pgio.AppendUint32(newBuf, uint32(len(fieldBytes)))
+			newBuf = append(newBuf, fieldBytes...)
+		} else {
+			newBuf = pgio.AppendInt32(newBuf, int32(-1))
+		}
+
+	}
+	return
+}
+
 // ROW allows deconstructing row values (records and composite types) into
 // fields directly without creating your own type and implementing decoder interfaces
 func ROW(isNull *bool, fields ...interface{}) BinaryDecoderFunc {