Fix #2385: Preserve multidimensional array dimensions in CopyFrom

copy_from_test.go

@@ -907,32 +907,123 @@ func TestConnCopyFromAutomaticStringConversion(t *testing.T) {
 func TestConnCopyFromAutomaticStringConversionArray(t *testing.T) {
 	t.Parallel()
 
+	for _, tt := range []struct {
+		name             string
+		multidimensional bool
+		inputRows        [][]any
+		want             [][]int64
+	}{
+		{
+			name:      "one-dimensional",
+			inputRows: [][]any{{[]string{"42"}}, {[]string{"7"}}, {[]string{"8", "9"}}},
+			want:      [][]int64{{42}, {7}, {8, 9}},
+		},
+		{
+			name:             "multidimensional",
+			multidimensional: true,
+			inputRows:        [][]any{{[][]string{{"10", "11"}, {"12", "13"}}}},
+			want:             [][]int64{{10, 11, 12, 13}},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+			defer cancel()
+
+			conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+			defer closeConn(t, conn)
+
+			if tt.multidimensional {
+				pgxtest.SkipCockroachDB(t, conn, "Server does not support nested arrays")
+			}
+
+			mustExec(t, conn, `create temporary table foo(a numeric[])`)
+
+			copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, pgx.CopyFromRows(tt.inputRows))
+			require.NoError(t, err)
+			require.EqualValues(t, len(tt.inputRows), copyCount)
+
+			// Test reads as int64 and flattened array for simplicity.
+			rows, _ := conn.Query(ctx, "select * from foo")
+			nums, err := pgx.CollectRows(rows, pgx.RowTo[[]int64])
+			require.NoError(t, err)
+			require.Equal(t, tt.want, nums)
+
+			ensureConnValid(t, conn)
+		})
+	}
+}
+
+// https://github.com/jackc/pgx/issues/2385
+func TestConnCopyFromPreservesMultidimensionalArrayString(t *testing.T) {
+	t.Parallel()
+
 	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()
 
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
+	pgxtest.SkipCockroachDB(t, conn, "Server does not support nested arrays")
+
 	mustExec(t, conn, `create temporary table foo(
-		a numeric[]
+		id int primary key,
+		a text[],
+		b text[][]
 	)`)
 
-	inputRows := [][]any{
-		{[]string{"42"}},
-		{[]string{"7"}},
-		{[]string{"8", "9"}},
-		{[][]string{{"10", "11"}, {"12", "13"}}},
+	// Reference row inserted with Exec, which sends the array literal in text format.
+	_, err := conn.Exec(ctx, `insert into foo (id, a, b) values (1, '{x,NULL,y}', '{{1,2},{3,NULL}}')`)
+	require.NoError(t, err)
+
+	// The same values as strings inserted with CopyFrom, which sends them in binary format.
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"id", "a", "b"}, pgx.CopyFromRows([][]any{
+		{2, "{x,NULL,y}", "{{1,2},{3,NULL}}"},
+		{3, nil, nil},
+	}))
+	require.NoError(t, err)
+	require.EqualValues(t, 2, copyCount)
+
+	// CopyFrom must store the same dimensions and elements as Exec, including NULLs.
+	expectedA := pgtype.Array[any]{
+		Elements: []any{"x", nil, "y"},
+		Dims:     []pgtype.ArrayDimension{{Length: 3, LowerBound: 1}},
+		Valid:    true,
+	}
+	expectedB := pgtype.Array[any]{
+		Elements: []any{"1", "2", "3", nil},
+		Dims: []pgtype.ArrayDimension{
+			{Length: 2, LowerBound: 1},
+			{Length: 2, LowerBound: 1},
+		},
+		Valid: true,
+	}
+
+	for _, id := range []int{1, 2} {
+		var gotA, gotB pgtype.Array[any]
+		err = conn.QueryRow(ctx, `select a, b from foo where id = $1`, id).Scan(&gotA, &gotB)
+		require.NoError(t, err)
+		require.Equal(t, expectedA, gotA, "row %d column a", id)
+		require.Equal(t, expectedB, gotB, "row %d column b", id)
 	}
 
-	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, pgx.CopyFromRows(inputRows))
+	// Dimensions are also visible through the PostgreSQL array inspection functions.
+	var ndims int
+	var dims string
+	var length int
+	err = conn.QueryRow(ctx, `select array_ndims(b), array_dims(b), array_length(b, 1) from foo where id = 2`).Scan(&ndims, &dims, &length)
 	require.NoError(t, err)
-	require.EqualValues(t, len(inputRows), copyCount)
+	require.Equal(t, 2, ndims)
+	require.Equal(t, "[1:2][1:2]", dims)
+	require.Equal(t, 2, length)
 
-	// Test reads as int64 and flattened array for simplicity.
-	rows, _ := conn.Query(ctx, "select * from foo")
-	nums, err := pgx.CollectRows(rows, pgx.RowTo[[]int64])
+	// NULL values stay NULL.
+	var aNull, bNull bool
+	err = conn.QueryRow(ctx, `select a is null, b is null from foo where id = 3`).Scan(&aNull, &bNull)
 	require.NoError(t, err)
-	require.Equal(t, [][]int64{{42}, {7}, {8, 9}, {10, 11, 12, 13}}, nums)
+	require.True(t, aNull)
+	require.True(t, bNull)
 
 	ensureConnValid(t, conn)
 }

values.go

@@ -53,6 +53,17 @@ func tryScanStringCopyValueThenEncode(m *pgtype.Map, buf []byte, oid uint32, arg
 		s = string(textBuf)
 	}
 
+	// Check if this is an array type. If so, use an ArraySetter to preserve
+	// array dimensions during CopyFrom.
+	if typ, ok := m.TypeForOID(oid); ok {
+		if _, ok := typ.Codec.(*pgtype.ArrayCodec); ok {
+			var arr pgtype.Array[any]
+			if err := m.Scan(oid, TextFormatCode, []byte(s), &arr); err == nil {
+				return m.Encode(oid, BinaryFormatCode, arr, buf)
+			}
+		}
+	}
+
 	var v any
 	err := m.Scan(oid, TextFormatCode, []byte(s), &v)
 	if err != nil {

values_internal_test.go

@@ -0,0 +1,93 @@
+package pgx
+
+import (
+	"encoding/binary"
+	"testing"
+
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/stretchr/testify/require"
+)
+
+// arrayBinaryHeader returns ndim, flags and the (length, lower bound) pairs from a binary array.
+func arrayBinaryHeader(t *testing.T, buf []byte) (uint32, uint32, [][2]int32) {
+	t.Helper()
+	require.GreaterOrEqual(t, len(buf), 12)
+	ndim := binary.BigEndian.Uint32(buf[0:4])
+	flags := binary.BigEndian.Uint32(buf[4:8])
+	var dims [][2]int32
+	for i := uint32(0); i < ndim; i++ {
+		off := 12 + i*8
+		require.GreaterOrEqual(t, len(buf), int(off)+8)
+		dims = append(dims, [2]int32{
+			int32(binary.BigEndian.Uint32(buf[off : off+4])),
+			int32(binary.BigEndian.Uint32(buf[off+4 : off+8])),
+		})
+	}
+	return ndim, flags, dims
+}
+
+// encodeCopyValue must preserve the dimensions of a multidimensional array given as a string.
+// https://github.com/jackc/pgx/issues/2385
+func TestEncodeCopyValuePreservesArrayDimensions(t *testing.T) {
+	t.Parallel()
+
+	m := pgtype.NewMap()
+	const textArrayOID = uint32(1009) // text[]
+	const int4OID = uint32(23)        // int4
+
+	// encode returns the binary representation of a single copy-format value (length prefix removed).
+	encode := func(t *testing.T, oid uint32, value any) []byte {
+		t.Helper()
+		buf, err := encodeCopyValue(m, nil, oid, value)
+		require.NoError(t, err)
+		require.GreaterOrEqual(t, len(buf), 4)
+		length := int32(binary.BigEndian.Uint32(buf))
+		require.EqualValues(t, len(buf)-4, length)
+		return buf[4:]
+	}
+
+	t.Run("two-dimensional array string", func(t *testing.T) {
+		ndim, flags, dims := arrayBinaryHeader(t, encode(t, textArrayOID, "{{1,2},{3,4}}"))
+		require.EqualValues(t, 2, ndim)
+		require.EqualValues(t, 0, flags)
+		require.Equal(t, [][2]int32{{2, 1}, {2, 1}}, dims)
+	})
+
+	t.Run("one-dimensional array string", func(t *testing.T) {
+		ndim, _, dims := arrayBinaryHeader(t, encode(t, textArrayOID, "{a,b,c}"))
+		require.EqualValues(t, 1, ndim)
+		require.Equal(t, [][2]int32{{3, 1}}, dims)
+	})
+
+	t.Run("array string with NULL element", func(t *testing.T) {
+		binaryValue := encode(t, textArrayOID, "{a,NULL,b}")
+		ndim, flags, dims := arrayBinaryHeader(t, binaryValue)
+		require.EqualValues(t, 1, ndim)
+		require.EqualValues(t, 1, flags) // ARRAY_HASNULL
+		require.Equal(t, [][2]int32{{3, 1}}, dims)
+
+		// Elements follow the header: int32 length (-1 for NULL) then the bytes.
+		off := 12 + len(dims)*8
+		elemLen := int32(binary.BigEndian.Uint32(binaryValue[off : off+4]))
+		require.EqualValues(t, 1, elemLen)
+		require.Equal(t, "a", string(binaryValue[off+4:off+4+int(elemLen)]))
+		off += 4 + int(elemLen)
+		elemLen = int32(binary.BigEndian.Uint32(binaryValue[off : off+4]))
+		require.EqualValues(t, -1, elemLen) // NULL
+		off += 4
+		elemLen = int32(binary.BigEndian.Uint32(binaryValue[off : off+4]))
+		require.EqualValues(t, 1, elemLen)
+		require.Equal(t, "b", string(binaryValue[off+4:off+4+int(elemLen)]))
+	})
+
+	t.Run("invalid array string still errors", func(t *testing.T) {
+		_, err := encodeCopyValue(m, nil, textArrayOID, "not an array")
+		require.Error(t, err)
+	})
+
+	t.Run("non-array type unaffected", func(t *testing.T) {
+		binaryValue := encode(t, int4OID, "42")
+		require.Len(t, binaryValue, 4)
+		require.EqualValues(t, 42, int32(binary.BigEndian.Uint32(binaryValue)))
+	})
+}