feat(pg): Use pgx batch to improve performance (#2604)

yevgenypats · web-flow · commit ed16cef82213 · 2022-10-11T08:13:56.000Z
Adds batching to improve performance over network + exposing this as configuration. Closes #2462
diff --git a/plugins/destination/postgresql/README.md b/plugins/destination/postgresql/README.md
@@ -19,3 +19,7 @@ This is the top level spec used by the PostgreSQL destination Plugin.
 
   Available: "error", "warn", "info", "debug", "trace"
   define if and in which level to log [`pgx`](https://github.com/jackc/pgx) call.
+
+- `batch_size` (int) (optional, defaults to 1000)
+
+  Number of rows to insert in a single batch.
diff --git a/plugins/destination/postgresql/client/client.go b/plugins/destination/postgresql/client/client.go
@@ -15,18 +15,15 @@ import (
 	pgxUUID "github.com/vgarvardt/pgx-google-uuid/v4"
 )
 
-type Spec struct {
-	ConnectionString string   `json:"connection_string,omitempty"`
-	PgxLogLevel      LogLevel `json:"pgx_log_level,omitempty"`
-}
-
 type Client struct {
 	conn                *pgxpool.Pool
 	logger              zerolog.Logger
 	spec                specs.Destination
 	currentDatabaseName string
 	currentSchemaName   string
 	pgType              pgType
+	batchSize           int
+	batch               *pgx.Batch
 }
 
 type pgTablePrimaryKeys struct {
@@ -73,12 +70,15 @@ const (
 func New(ctx context.Context, logger zerolog.Logger, spec specs.Destination) (plugins.DestinationClient, error) {
 	c := &Client{
 		logger: logger.With().Str("module", "pg-dest").Logger(),
+		batch:  &pgx.Batch{},
 	}
 	var specPostgreSql Spec
 	c.spec = spec
 	if err := spec.UnmarshalSpec(&specPostgreSql); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal postgresql spec: %w", err)
 	}
+	specPostgreSql.SetDefaults()
+	c.batchSize = specPostgreSql.BatchSize
 
 	logLevel, err := pgx.LogLevelFromString(specPostgreSql.PgxLogLevel.String())
 	if err != nil {
@@ -117,14 +117,20 @@ func New(ctx context.Context, logger zerolog.Logger, spec specs.Destination) (pl
 }
 
 func (c *Client) Close(ctx context.Context) error {
+	var err error
 	if c.conn == nil {
 		return fmt.Errorf("client already closed or not initialized")
 	}
+	if c.batch.Len() > 0 {
+		br := c.conn.SendBatch(ctx, c.batch)
+		err = br.Close()
+		c.batch = &pgx.Batch{}
+	}
 	if c.conn != nil {
 		c.conn.Close()
 		c.conn = nil
 	}
-	return nil
+	return err
 }
 
 func (c *Client) currentDatabase(ctx context.Context) (string, error) {
diff --git a/plugins/destination/postgresql/client/deletestale.go b/plugins/destination/postgresql/client/deletestale.go
@@ -9,7 +9,7 @@ import (
 	"github.com/jackc/pgx/v4"
 )
 
-func (p *Client) DeleteStale(ctx context.Context, table string, source string, syncTime time.Time) error {
+func (c *Client) DeleteStale(ctx context.Context, table string, source string, syncTime time.Time) error {
 	var sb strings.Builder
 	sb.WriteString("delete from ")
 	sb.WriteString(pgx.Identifier{table}.Sanitize())
@@ -18,7 +18,7 @@ func (p *Client) DeleteStale(ctx context.Context, table string, source string, s
 	sb.WriteString(" = $1 and ")
 	sb.WriteString(schema.CqSyncTimeColumn.Name)
 	sb.WriteString(" < $2")
-	if _, err := p.conn.Exec(ctx, sb.String(), source, syncTime); err != nil {
+	if _, err := c.conn.Exec(ctx, sb.String(), source, syncTime); err != nil {
 		return err
 	}
 	return nil
diff --git a/plugins/destination/postgresql/client/deletestale_test.go b/plugins/destination/postgresql/client/deletestale_test.go
@@ -15,6 +15,7 @@ func TestDeleteStale(t *testing.T) {
 		WriteMode: specs.WriteModeOverwriteDeleteStale,
 		Spec: &Spec{
 			ConnectionString: getTestConnection(),
+			BatchSize:        1,
 		},
 	})
 	if err != nil {
diff --git a/plugins/destination/postgresql/client/migrate_test.go b/plugins/destination/postgresql/client/migrate_test.go
@@ -15,6 +15,7 @@ func TestMigrate(t *testing.T) {
 		WriteMode: specs.WriteModeOverwriteDeleteStale,
 		Spec: &Spec{
 			ConnectionString: getTestConnection(),
+			BatchSize:        1,
 		},
 	})
 	if err != nil {
diff --git a/plugins/destination/postgresql/client/spec.go b/plugins/destination/postgresql/client/spec.go
@@ -0,0 +1,15 @@
+package client
+
+type Spec struct {
+	ConnectionString string   `json:"connection_string,omitempty"`
+	PgxLogLevel      LogLevel `json:"pgx_log_level,omitempty"`
+	BatchSize        int      `json:"batch_size,omitempty"`
+}
+
+const defaultBatchSize = 1000
+
+func (s *Spec) SetDefaults() {
+	if s.BatchSize <= 0 {
+		s.BatchSize = defaultBatchSize
+	}
+}
diff --git a/plugins/destination/postgresql/client/write.go b/plugins/destination/postgresql/client/write.go
@@ -9,17 +9,23 @@ import (
 	"github.com/jackc/pgx/v4"
 )
 
-func (p *Client) Write(ctx context.Context, table string, data map[string]interface{}) error {
+func (c *Client) Write(ctx context.Context, table string, data map[string]interface{}) error {
 	var sql string
 	var values []interface{}
-	if p.spec.WriteMode == specs.WriteModeAppend {
+
+	if c.spec.WriteMode == specs.WriteModeAppend {
 		sql, values = insert(table, data)
 	} else {
 		sql, values = upsert(table, data)
 	}
-	_, err := p.conn.Exec(ctx, sql, values...)
-	if err != nil {
-		return fmt.Errorf("failed to insert data with sql '%s': %w", sql, err)
+	c.batch.Queue(sql, values...)
+	if c.batch.Len() >= c.batchSize {
+		br := c.conn.SendBatch(ctx, c.batch)
+		if err := br.Close(); err != nil {
+			c.batch = &pgx.Batch{}
+			return fmt.Errorf("failed to execute batch: %v", err)
+		}
+		c.batch = &pgx.Batch{}
 	}
 	return nil
 }
diff --git a/plugins/destination/postgresql/client/write_test.go b/plugins/destination/postgresql/client/write_test.go
@@ -18,6 +18,7 @@ func TestWriteOverwriteDeleteStale(t *testing.T) {
 		WriteMode: specs.WriteModeOverwriteDeleteStale,
 		Spec: &Spec{
 			ConnectionString: getTestConnection(),
+			BatchSize:        1,
 		},
 	})
 	if err != nil {
@@ -73,6 +74,7 @@ func TestWriteAppend(t *testing.T) {
 		WriteMode: specs.WriteModeAppend,
 		Spec: &Spec{
 			ConnectionString: getTestConnection(),
+			BatchSize:        1,
 		},
 	})
 	if err != nil {
