internal/postgres: support bulk inserts exceeding 65535 params

Postgres has a hard limit of 2^16-1 query parameters, which we were
exceeding for some modules (notably Kubernetes). This CL works around
this limitation by executing multiple queries.

Fixes b/133512292

Change-Id: I17087ea54ddab3515068f1ef131dd2ba87165a5f
Reviewed-on: https://team-review.git.corp.google.com/c/golang/discovery/+/476518
Reviewed-by: Julie Qiu <julieqiu@google.com>

1 files changed, 34 insertions, 0 deletions

diff --git a/internal/postgres/postgres_test.go b/internal/postgres/postgres_test.go
index 76acfee1..5b3d02df 100644
--- a/internal/postgres/postgres_test.go
+++ b/internal/postgres/postgres_test.go
@@ -190,6 +190,40 @@ func TestBulkInsert(t *testing.T) {
 	}
 }
 
+func TestLargeBulkInsert(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+	if _, err := testDB.ExecContext(ctx, `CREATE TEMPORARY TABLE test_large_bulk (i BIGINT);`); err != nil {
+		t.Fatal(err)
+	}
+	const size = 150000
+	vals := make([]interface{}, size)
+	for i := 0; i < size; i++ {
+		vals[i] = i + 1
+	}
+	if err := testDB.Transact(func(tx *sql.Tx) error {
+		return bulkInsert(ctx, tx, "test_large_bulk", []string{"i"}, vals, "")
+	}); err != nil {
+		t.Fatal(err)
+	}
+	rows, err := testDB.QueryContext(ctx, `SELECT i FROM test_large_bulk;`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sum := int64(0)
+	for rows.Next() {
+		var i int64
+		if err := rows.Scan(&i); err != nil {
+			t.Fatal(err)
+		}
+		sum += i
+	}
+	var want int64 = size * (size + 1) / 2
+	if sum != want {
+		t.Errorf("sum = %d, want %d", sum, want)
+	}
+}
+
 func TestPostgres_ReadAndWriteVersionAndPackages(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()