Support writes to Postgres with pgx driver (#28521)

apache · Sep 19, 2023 · 9c4e8bf · 9c4e8bf
1 parent c7274e9
commit 9c4e8bf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git 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
@@ -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,