From 9c4e8bff7c97f513f9df498debdbcc08698413d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johanna=20=C3=96jeling?=
 <51084516+johannaojeling@users.noreply.github.com>
Date: Tue, 19 Sep 2023 17:03:11 +0200
Subject: [PATCH] Support writes to Postgres with pgx driver (#28521)

---
 sdks/go/pkg/beam/io/databaseio/writer.go      | 2 +-
 sdks/go/pkg/beam/io/databaseio/writer_test.go | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/sdks/go/pkg/beam/io/databaseio/writer.go b/sdks/go/pkg/beam/io/databaseio/writer.go
index 4719831581ff..ebd805b080d7 100644
--- a/sdks/go/pkg/beam/io/databaseio/writer.go
+++ b/sdks/go/pkg/beam/io/databaseio/writer.go
@@ -108,7 +108,7 @@ type valueTemplateGenerator struct {
 
 func (v *valueTemplateGenerator) generate(rowCount int, columnColunt int) string {
 	switch v.driver {
-	case "postgres":
+	case "postgres", "pgx":
 		// the point is to generate ($1,$2),($3,$4)
 		valueTemplates := make([]string, rowCount)
 		for i := 0; i < rowCount; i++ {
diff --git a/sdks/go/pkg/beam/io/databaseio/writer_test.go b/sdks/go/pkg/beam/io/databaseio/writer_test.go
index 24f617a9b857..91d7fb9246dd 100644
--- a/sdks/go/pkg/beam/io/databaseio/writer_test.go
+++ b/sdks/go/pkg/beam/io/databaseio/writer_test.go
@@ -39,6 +39,12 @@ func TestValueTemplateGenerator_generate(t *testing.T) {
 			columnCount: 10,
 			expected:    "",
 		},
+		{
+			generator:   &valueTemplateGenerator{"pgx"},
+			rowCount:    4,
+			columnCount: 3,
+			expected:    "($1,$2,$3),($4,$5,$6),($7,$8,$9),($10,$11,$12)",
+		},
 		{
 			generator:   &valueTemplateGenerator{"mysql"},
 			rowCount:    4,