From f42af358847797376847ffa58d105d8373fbddaf Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Tue, 20 Dec 2022 20:11:04 -0600
Subject: [PATCH] Add support for single dimensional arrays

https://github.com/jackc/pgx/issues/1442
---
 pgtype/array_codec_test.go | 23 ++++++++++++++++
 pgtype/builtin_wrappers.go | 40 +++++++++++++++++++++++++++
 pgtype/pgtype.go           | 56 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 119 insertions(+)

diff --git a/pgtype/array_codec_test.go b/pgtype/array_codec_test.go
index a558d0fc..ede104ac 100644
--- a/pgtype/array_codec_test.go
+++ b/pgtype/array_codec_test.go
@@ -126,6 +126,29 @@ func TestArrayCodecAnySlice(t *testing.T) {
 	})
 }
 
+// https://github.com/jackc/pgx/issues/1442
+func TestArrayCodecAnyArray(t *testing.T) {
+	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		type _point3 [3]float32
+
+		for i, tt := range []struct {
+			expected any
+		}{
+			{_point3{0, 0, 0}},
+			{_point3{1, 2, 3}},
+		} {
+			var actual _point3
+			err := conn.QueryRow(
+				ctx,
+				"select $1::float4[]",
+				tt.expected,
+			).Scan(&actual)
+			assert.NoErrorf(t, err, "%d", i)
+			assert.Equalf(t, tt.expected, actual, "%d", i)
+		}
+	})
+}
+
 // https://github.com/jackc/pgx/issues/1273#issuecomment-1218262703
 func TestArrayCodecSliceArgConversion(t *testing.T) {
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
diff --git a/pgtype/builtin_wrappers.go b/pgtype/builtin_wrappers.go
index e34dd578..8bf367c1 100644
--- a/pgtype/builtin_wrappers.go
+++ b/pgtype/builtin_wrappers.go
@@ -910,3 +910,43 @@ func (a *anyMultiDimSliceArray) ScanIndexType() any {
 	}
 	return reflect.New(lowestSliceType.Elem()).Interface()
 }
+
+type anyArrayArrayReflect struct {
+	array reflect.Value
+}
+
+func (a anyArrayArrayReflect) Dimensions() []ArrayDimension {
+	return []ArrayDimension{{Length: int32(a.array.Len()), LowerBound: 1}}
+}
+
+func (a anyArrayArrayReflect) Index(i int) any {
+	return a.array.Index(i).Interface()
+}
+
+func (a anyArrayArrayReflect) IndexType() any {
+	return reflect.New(a.array.Type().Elem()).Elem().Interface()
+}
+
+func (a *anyArrayArrayReflect) SetDimensions(dimensions []ArrayDimension) error {
+	if dimensions == nil {
+		return fmt.Errorf("anyArrayArrayReflect: cannot scan NULL into %v", a.array.Type().String())
+	}
+
+	if len(dimensions) != 1 {
+		return fmt.Errorf("anyArrayArrayReflect: cannot scan multi-dimensional array into %v", a.array.Type().String())
+	}
+
+	if int(dimensions[0].Length) != a.array.Len() {
+		return fmt.Errorf("anyArrayArrayReflect: cannot scan array with length %v into %v", dimensions[0].Length, a.array.Type().String())
+	}
+
+	return nil
+}
+
+func (a *anyArrayArrayReflect) ScanIndex(i int) any {
+	return a.array.Index(i).Addr().Interface()
+}
+
+func (a *anyArrayArrayReflect) ScanIndexType() any {
+	return reflect.New(a.array.Type().Elem()).Interface()
+}
diff --git a/pgtype/pgtype.go b/pgtype/pgtype.go
index 2aa96e68..7cfa3d58 100644
--- a/pgtype/pgtype.go
+++ b/pgtype/pgtype.go
@@ -223,6 +223,7 @@ func NewMap() *Map {
 			TryWrapStructEncodePlan,
 			TryWrapSliceEncodePlan,
 			TryWrapMultiDimSliceEncodePlan,
+			TryWrapArrayEncodePlan,
 		},
 
 		TryWrapScanPlanFuncs: []TryWrapScanPlanFunc{
@@ -232,6 +233,7 @@ func NewMap() *Map {
 			TryWrapStructScanPlan,
 			TryWrapPtrSliceScanPlan,
 			TryWrapPtrMultiDimSliceScanPlan,
+			TryWrapPtrArrayScanPlan,
 		},
 	}
 
@@ -1139,6 +1141,31 @@ func (plan *wrapPtrMultiDimSliceScanPlan) Scan(src []byte, target any) error {
 	return plan.next.Scan(src, &anyMultiDimSliceArray{slice: reflect.ValueOf(target).Elem()})
 }
 
+// TryWrapPtrArrayScanPlan tries to wrap a pointer to a single dimension array.
+func TryWrapPtrArrayScanPlan(target any) (plan WrappedScanPlanNextSetter, nextValue any, ok bool) {
+	targetValue := reflect.ValueOf(target)
+	if targetValue.Kind() != reflect.Ptr {
+		return nil, nil, false
+	}
+
+	targetElemValue := targetValue.Elem()
+
+	if targetElemValue.Kind() == reflect.Array {
+		return &wrapPtrArrayReflectScanPlan{}, &anyArrayArrayReflect{array: targetElemValue}, true
+	}
+	return nil, nil, false
+}
+
+type wrapPtrArrayReflectScanPlan struct {
+	next ScanPlan
+}
+
+func (plan *wrapPtrArrayReflectScanPlan) SetNext(next ScanPlan) { plan.next = next }
+
+func (plan *wrapPtrArrayReflectScanPlan) Scan(src []byte, target any) error {
+	return plan.next.Scan(src, &anyArrayArrayReflect{array: reflect.ValueOf(target).Elem()})
+}
+
 // PlanScan prepares a plan to scan a value into target.
 func (m *Map) PlanScan(oid uint32, formatCode int16, target any) ScanPlan {
 	oidMemo := m.memoizedScanPlans[oid]
@@ -1941,6 +1968,35 @@ func (plan *wrapMultiDimSliceEncodePlan) Encode(value any, buf []byte) (newBuf [
 	return plan.next.Encode(&w, buf)
 }
 
+func TryWrapArrayEncodePlan(value any) (plan WrappedEncodePlanNextSetter, nextValue any, ok bool) {
+	if _, ok := value.(driver.Valuer); ok {
+		return nil, nil, false
+	}
+
+	if valueType := reflect.TypeOf(value); valueType != nil && valueType.Kind() == reflect.Array {
+		w := anyArrayArrayReflect{
+			array: reflect.ValueOf(value),
+		}
+		return &wrapArrayEncodeReflectPlan{}, w, true
+	}
+
+	return nil, nil, false
+}
+
+type wrapArrayEncodeReflectPlan struct {
+	next EncodePlan
+}
+
+func (plan *wrapArrayEncodeReflectPlan) SetNext(next EncodePlan) { plan.next = next }
+
+func (plan *wrapArrayEncodeReflectPlan) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	w := anyArrayArrayReflect{
+		array: reflect.ValueOf(value),
+	}
+
+	return plan.next.Encode(w, buf)
+}
+
 func newEncodeError(value any, m *Map, oid uint32, formatCode int16, err error) error {
 	var format string
 	switch formatCode {