From 45100245a801780cbae20599191559dcfbe524bd Mon Sep 17 00:00:00 2001
From: Peter Dotchev <dotchev@gmail.com>
Date: Wed, 15 Apr 2020 13:27:09 +0300
Subject: [PATCH] Insert labels in batches to avoid pq param limit (#491)

Cherry-pick PR #489
---
 storage/postgres/storage.go                   | 29 +++++++++++++++++--
 test/common/application.yml                   |  2 +-
 test/configuration_test/configuration_test.go |  2 +-
 test/visibility_test/visibility_test.go       | 18 ++++++++++++
 4 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/storage/postgres/storage.go b/storage/postgres/storage.go
index f9e4f353e..25154b9a3 100644
--- a/storage/postgres/storage.go
+++ b/storage/postgres/storage.go
@@ -217,9 +217,32 @@ func (ps *Storage) createLabels(ctx context.Context, labels []PostgresLabel) err
 		strings.Join(dbTags, ", :"),
 	)
 
-	log.C(ctx).Debugf("Executing query %s", sqlQuery)
-	_, err := ps.pgDB.NamedExecContext(ctx, sqlQuery, labels)
-	return checkIntegrityViolation(ctx, err)
+	// break into batches so the PostgreSQL limit of 65535 parameters is not exceeded
+	const maxParams = 65535
+	maxRows := maxParams / len(dbTags)
+	for len(labels) > 0 {
+		rows := min(len(labels), maxRows)
+		log.C(ctx).Debugf("Executing query %s", sqlQuery)
+		result, err := ps.pgDB.NamedExecContext(ctx, sqlQuery, labels[:rows])
+		if err != nil {
+			return checkIntegrityViolation(ctx, err)
+		}
+		rowsAffected, err := result.RowsAffected()
+		if err != nil {
+			log.C(ctx).Debugf("Could not get number of affected rows: %v", err)
+		} else {
+			log.C(ctx).Debugf("%d rows inserted", rowsAffected)
+		}
+		labels = labels[rows:]
+	}
+	return nil
+}
+
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
 }
 
 func (ps *Storage) deleteLabels(ctx context.Context, objectType types.ObjectType, entityID string, removedLabels types.Labels) error {
diff --git a/test/common/application.yml b/test/common/application.yml
index 0b0d1d7ed..5a54369aa 100644
--- a/test/common/application.yml
+++ b/test/common/application.yml
@@ -1,5 +1,5 @@
 server:
-  request_timeout: 4000ms
+  request_timeout: 30s
   shutdown_timeout: 4000ms
   port: 1234
 httpclient:
diff --git a/test/configuration_test/configuration_test.go b/test/configuration_test/configuration_test.go
index 140ca7724..ba8e10c69 100644
--- a/test/configuration_test/configuration_test.go
+++ b/test/configuration_test/configuration_test.go
@@ -97,7 +97,7 @@ var _ = Describe("Service Manager Config API", func() {
 					"max_body_bytes": 1048576,
 					"max_header_bytes": 1024,
 					"port": 1234,
-					"request_timeout": "4000ms",
+					"request_timeout": "30s",
 					"shutdown_timeout": "4000ms"
 				},
 				"storage": {
diff --git a/test/visibility_test/visibility_test.go b/test/visibility_test/visibility_test.go
index 6345ac8c5..9d2f6a90e 100644
--- a/test/visibility_test/visibility_test.go
+++ b/test/visibility_test/visibility_test.go
@@ -250,6 +250,24 @@ var _ = test.DescribeTestsFor(test.TestCase{
 						})
 					})
 
+					Context("When many labels are provided", func() {
+						It("should return 201", func() {
+							// see https://github.com/lib/pq/blob/master/conn.go#L1282
+							const labelCount = 20000 // 20000 * 6 > 65535 - max postgres parameter number
+							orgs := make(common.Array, labelCount)
+							for i := range orgs {
+								orgs[i] = fmt.Sprintf("org-id-%d", i)
+							}
+							postVisibilityRequestWithLabels["labels"] = common.Object{
+								"org_id": orgs,
+							}
+							ctx.SMWithOAuth.POST(web.VisibilitiesURL).
+								WithJSON(postVisibilityRequestWithLabels).
+								Expect().Status(http.StatusCreated).
+								JSON().Object().Path("$.labels.org_id").Array().ContainsOnly(orgs...)
+						})
+					})
+
 					Context("When creating labeled visibility for which a public one exists", func() {
 						It("Should return 409", func() {
 							ctx.SMWithOAuth.POST(web.VisibilitiesURL).