From ee001a7caebc75cc51d0068e9109b99d8a20755d Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Mon, 1 May 2017 19:46:37 -0500
Subject: [PATCH] Fix queries with more than 32 columns

fixes #270
---
 pgproto3/data_row.go |  6 ++++-
 query_test.go        | 53 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/pgproto3/data_row.go b/pgproto3/data_row.go
index 6b27f728..3e600e84 100644
--- a/pgproto3/data_row.go
+++ b/pgproto3/data_row.go
@@ -25,7 +25,11 @@ func (dst *DataRow) Decode(src []byte) error {
 	// large reallocate. This is too avoid one row with many columns from
 	// permanently allocating memory.
 	if cap(dst.Values) < fieldCount || cap(dst.Values)-fieldCount > 32 {
-		dst.Values = make([][]byte, fieldCount, 32)
+		newCap := 32
+		if newCap < fieldCount {
+			newCap = fieldCount
+		}
+		dst.Values = make([][]byte, fieldCount, newCap)
 	} else {
 		dst.Values = dst.Values[:fieldCount]
 	}
diff --git a/query_test.go b/query_test.go
index c1ca480a..801b34dd 100644
--- a/query_test.go
+++ b/query_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -46,6 +47,58 @@ func TestConnQueryScan(t *testing.T) {
 	}
 }
 
+func TestConnQueryScanWithManyColumns(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnect(t, *defaultConnConfig)
+	defer closeConn(t, conn)
+
+	columnCount := 1000
+	sql := "select "
+	for i := 0; i < columnCount; i++ {
+		if i > 0 {
+			sql += ","
+		}
+		sql += fmt.Sprintf(" %d", i)
+	}
+	sql += " from generate_series(1,5)"
+
+	dest := make([]int, columnCount)
+
+	var rowCount int
+
+	rows, err := conn.Query(sql)
+	if err != nil {
+		t.Fatalf("conn.Query failed: %v", err)
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		destPtrs := make([]interface{}, columnCount)
+		for i := range destPtrs {
+			destPtrs[i] = &dest[i]
+		}
+		if err := rows.Scan(destPtrs...); err != nil {
+			t.Fatalf("rows.Scan failed: %v", err)
+		}
+		rowCount++
+
+		for i := range dest {
+			if dest[i] != i {
+				t.Errorf("dest[%d] => %d, want %d", i, dest[i], i)
+			}
+		}
+	}
+
+	if rows.Err() != nil {
+		t.Fatalf("conn.Query failed: %v", err)
+	}
+
+	if rowCount != 5 {
+		t.Errorf("rowCount => %d, want %d", rowCount, 5)
+	}
+}
+
 func TestConnQueryValues(t *testing.T) {
 	t.Parallel()