Repository: xataio/pgstream
Pull request #125 (2025-01-24 14:46:40)

Add Postgres output with DML support
This PR adds support for a PostgreSQL pgstream output for DML events (DDL events implemented in upcoming PRs). The postgres batch writer reuses the batch sender, parsing the DML wal events from the replication slot on the source postgres database. 

The relevant configuration has been added for now, but the readme will be updated once the output is completely implemented.

The docker compose file is updated to add another postgres database, and profiles so that only the relevant resources are started.

github-actions[bot]: ### Merging this branch will **increase** overall coverage

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/cmd | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/stream | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/wal/processor/postgres | 81.25% (**+81.25%**) | :star2: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/cmd/config.go | 0.00% (ø) | 71 (+4) | 0 | 71 (+4) |  |
| github.com/xataio/pgstream/pkg/stream/config.go | 0.00% (ø) | 10 | 0 | 10 |  |
| github.com/xataio/pgstream/pkg/stream/stream_run.go | 0.00% (ø) | 273 (+5) | 0 | 273 (+5) |  |
| github.com/xataio/pgstream/pkg/wal/processor/postgres/config.go | 0.00% (ø) | 0 | 0 | 0 |  |
| github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_batch_writer.go | 53.85% (**+53.85%**) | 39 (+39) | 21 (+21) | 18 (+18) | :star2: |
| github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_query_msg.go | 100.00% (**+100.00%**) | 2 (+2) | 2 (+2) | 0 | :star2: |
| github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_wal_adapter.go | 100.00% (**+100.00%**) | 55 (+55) | 55 (+55) | 0 | :star2: |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/pkg/wal/processor/postgres/helper_test.go
- github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_batch_writer_test.go
- github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_wal_adapter_test.go

</details>


Patch:
diff --git a/build/docker/docker-compose.yml b/build/docker/docker-compose.yml
index ea3d759..5ab0cae 100644
--- a/build/docker/docker-compose.yml
+++ b/build/docker/docker-compose.yml
@@ -1,23 +1,37 @@
-version: '3.8'
+version: "3.8"
 services:
-  db:
+  db1:
     build: postgres
     restart: always
     environment:
       - POSTGRES_USER=postgres
       - POSTGRES_PASSWORD=postgres
     ports:
-      - '5432:5432'
+      - "5432:5432"
     volumes:
       - data:/var/lib/postgresql/data
       - ./postgres/postgresql.conf:/etc/postgresql.conf
     command: -c config_file=/etc/postgresql.conf
+    profiles: ["pg2os", "pg2webhook", "pg2pg"]
+
+  db2:
+    build: postgres
+    restart: always
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+    ports:
+      - "7654:5432"
+    volumes:
+      - ./postgres/postgresql.conf:/etc/postgresql.conf
+    command: -c config_file=/etc/postgresql.conf
+    profiles: ["pg2pg"]
 
   kafka:
-    image: 'bitnami/kafka:latest'
+    image: "bitnami/kafka:latest"
     ports:
-     - '9092:9092'
-     - '9094:9094'
+      - "9092:9092"
+      - "9094:9094"
     environment:
       - KAFKA_CFG_NODE_ID=0
       - KAFKA_CFG_PROCESS_ROLES=controller,broker
@@ -26,6 +40,7 @@ services:
       - KAFKA_CFG_CONTROLLER_QUORUM_VOTERS=0@kafka:9093
       - KAFKA_CFG_CONTROLLER_LISTENER_NAMES=CONTROLLER
       - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://localhost:9092,EXTERNAL://localhost:9094
+    profiles: ["kafka"]
 
   opensearch:
     image: "opensearchproject/opensearch:2.7.0"
@@ -43,10 +58,12 @@ services:
       memlock:
         soft: -1
         hard: -1
+    profiles: ["pg2os"]
   opensearch-dashboards:
     image: "opensearchproject/opensearch-dashboards:2.7.0"
     ports: ["5601:5601"]
     environment: ["OPENSEARCH_HOSTS=http://opensearch:9200"] # 9201 to access externally, 9200 internally
+    profiles: ["pg2os"]
 
 volumes:
   data:
diff --git a/cmd/config.go b/cmd/config.go
index 31e1869..8d8e7ce 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -17,6 +17,7 @@ import (
 	"github.com/xataio/pgstream/pkg/wal/processor/batch"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
+	"github.com/xataio/pgstream/pkg/wal/processor/postgres"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
 	"github.com/xataio/pgstream/pkg/wal/processor/search/store"
 	"github.com/xataio/pgstream/pkg/wal/processor/webhook/notifier"
@@ -143,6 +144,7 @@ func parseProcessorConfig() stream.ProcessorConfig {
 		Kafka:    parseKafkaProcessorConfig(),
 		Search:   parseSearchProcessorConfig(),
 		Webhook:  parseWebhookProcessorConfig(),
+		Postgres: parsePostgresProcessorConfig(),
 		Injector: parseInjectorConfig(),
 	}
 }
@@ -231,6 +233,25 @@ func parseWebhookProcessorConfig() *stream.WebhookProcessorConfig {
 	}
 }
 
+func parsePostgresProcessorConfig() *stream.PostgresProcessorConfig {
+	targetPostgresURL := viper.GetString("PGSTREAM_POSTGRES_WRITER_TARGET_URL")
+	if targetPostgresURL == "" {
+		return nil
+	}
+
+	return &stream.PostgresProcessorConfig{
+		BatchWriter: postgres.Config{
+			URL: targetPostgresURL,
+			BatchConfig: batch.Config{
+				BatchTimeout:  viper.GetDuration("PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT"),
+				MaxBatchBytes: viper.GetInt64("PGSTREAM_POSTGRES_WRITER_BATCH_BYTES"),
+				MaxBatchSize:  viper.GetInt64("PGSTREAM_POSTGRES_WRITER_BATCH_SIZE"),
+				MaxQueueBytes: viper.GetInt64("PGSTREAM_POSTGRES_WRITER_MAX_QUEUE_BYTES"),
+			},
+		},
+	}
+}
+
 func parseBackoffConfig(prefix string) backoff.Config {
 	return backoff.Config{
 		Exponential: parseExponentialBackoffConfig(prefix),
diff --git a/pg2pg.env b/pg2pg.env
new file mode 100644
index 0000000..1c9192c
--- /dev/null
+++ b/pg2pg.env
@@ -0,0 +1,13 @@
+# Listener config
+PGSTREAM_POSTGRES_LISTENER_URL="postgres://postgres:postgres@localhost?sslmode=disable"
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLES="*"
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_SCHEMA_WORKERS=4
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS=4
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS=1
+
+# Processor config
+PGSTREAM_INJECTOR_STORE_POSTGRES_URL="postgres://postgres:postgres@localhost?sslmode=disable"
+PGSTREAM_POSTGRES_WRITER_TARGET_URL="postgres://postgres:postgres@localhost:7654?sslmode=disable"
+PGSTREAM_POSTGRES_WRITER_BATCH_SIZE=100
+PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT=5s
diff --git a/pkg/stream/config.go b/pkg/stream/config.go
index 74d76cd..f6466ef 100644
--- a/pkg/stream/config.go
+++ b/pkg/stream/config.go
@@ -11,6 +11,7 @@ import (
 	pglistener "github.com/xataio/pgstream/pkg/wal/listener/postgres"
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
+	"github.com/xataio/pgstream/pkg/wal/processor/postgres"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
 	"github.com/xataio/pgstream/pkg/wal/processor/search/store"
 	"github.com/xataio/pgstream/pkg/wal/processor/webhook/notifier"
@@ -42,6 +43,7 @@ type ProcessorConfig struct {
 	Kafka    *KafkaProcessorConfig
 	Search   *SearchProcessorConfig
 	Webhook  *WebhookProcessorConfig
+	Postgres *PostgresProcessorConfig
 	Injector *injector.Config
 }
 
@@ -61,6 +63,10 @@ type WebhookProcessorConfig struct {
 	SubscriptionStore  WebhookSubscriptionStoreConfig
 }
 
+type PostgresProcessorConfig struct {
+	BatchWriter postgres.Config
+}
+
 type WebhookSubscriptionStoreConfig struct {
 	URL                  string
 	CacheEnabled         bool
@@ -72,7 +78,7 @@ func (c *Config) IsValid() error {
 		return errors.New("need at least one listener configured")
 	}
 
-	if c.Processor.Kafka == nil && c.Processor.Search == nil && c.Processor.Webhook == nil {
+	if c.Processor.Kafka == nil && c.Processor.Search == nil && c.Processor.Webhook == nil && c.Processor.Postgres == nil {
 		return errors.New("need at least one processor configured")
 	}
 
diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go
index 1e37359..fca3d2e 100644
--- a/pkg/stream/stream_run.go
+++ b/pkg/stream/stream_run.go
@@ -22,6 +22,7 @@ import (
 	"github.com/xataio/pgstream/pkg/wal/processor/injector"
 	processinstrumentation "github.com/xataio/pgstream/pkg/wal/processor/instrumentation"
 	kafkaprocessor "github.com/xataio/pgstream/pkg/wal/processor/kafka"
+	pgwriter "github.com/xataio/pgstream/pkg/wal/processor/postgres"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
 	searchinstrumentation "github.com/xataio/pgstream/pkg/wal/processor/search/instrumentation"
 	"github.com/xataio/pgstream/pkg/wal/processor/search/store"
@@ -199,6 +200,18 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, instrumentat
 			return notifier.Notify(ctx)
 		})
 
+	case config.Processor.Postgres != nil:
+		pgBatchWriter, err := pgwriter.NewBatchWriter(ctx,
+			&config.Processor.Postgres.BatchWriter,
+			pgwriter.WithLogger(logger),
+			pgwriter.WithCheckpoint(checkpoint))
+		if err != nil {
+			return err
+		}
+		defer pgBatchWriter.Close()
+
+		processor = pgBatchWriter
+
 	default:
 		return errors.New("no processor found")
 	}
diff --git a/pkg/wal/processor/postgres/config.go b/pkg/wal/processor/postgres/config.go
new file mode 100644
index 0000000..54c8414
--- /dev/null
+++ b/pkg/wal/processor/postgres/config.go
@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"github.com/xataio/pgstream/pkg/wal/processor/batch"
+)
+
+type Config struct {
+	URL         string
+	BatchConfig batch.Config
+}
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer.go b/pkg/wal/processor/postgres/postgres_batch_writer.go
new file mode 100644
index 0000000..dbc533b
--- /dev/null
+++ b/pkg/wal/processor/postgres/postgres_batch_writer.go
@@ -0,0 +1,137 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/wal"
+	"github.com/xataio/pgstream/pkg/wal/checkpointer"
+	"github.com/xataio/pgstream/pkg/wal/processor"
+	"github.com/xataio/pgstream/pkg/wal/processor/batch"
+)
+
+// BatchWriter is a WAL processor implementation that batches and writes wal
+// events to a Postgres instance.
+type BatchWriter struct {
+	logger  loglib.Logger
+	pgConn  pglib.Querier
+	adapter walAdapter
+
+	batchSender queryBatchSender
+
+	// optional checkpointer callback to mark what was safely processed
+	checkpointer checkpointer.Checkpoint
+}
+
+type Option func(*BatchWriter)
+
+type queryBatchSender interface {
+	AddToBatch(context.Context, *batch.WALMessage[*query]) error
+	Close()
+	Send(context.Context) error
+}
+
+// NewBatchWriter returns a postgres processor that batches and writes data to
+// the configured postgres instance.
+func NewBatchWriter(ctx context.Context, config *Config, opts ...Option) (*BatchWriter, error) {
+	pgConn, err := pglib.NewConnPool(ctx, config.URL)
+	if err != nil {
+		return nil, err
+	}
+
+	w := &BatchWriter{
+		logger:  loglib.NewNoopLogger(),
+		pgConn:  pgConn,
+		adapter: &adapter{},
+	}
+
+	for _, opt := range opts {
+		opt(w)
+	}
+
+	w.batchSender, err = batch.NewSender(&config.BatchConfig, w.sendBatch, w.logger)
+	if err != nil {
+		return nil, err
+	}
+
+	// start the send process in the background
+	go func() {
+		if err := w.batchSender.Send(ctx); err != nil {
+			w.logger.Error(err, "sending stopped")
+		}
+	}()
+
+	return w, nil
+}
+
+func WithLogger(l loglib.Logger) Option {
+	return func(w *BatchWriter) {
+		w.logger = loglib.NewLogger(l).WithFields(loglib.Fields{
+			loglib.ModuleField: "postgres_batch_writer",
+		})
+	}
+}
+
+func WithCheckpoint(c checkpointer.Checkpoint) Option {
+	return func(w *BatchWriter) {
+		w.checkpointer = c
+	}
+}
+
+// ProcessWalEvent is called on every new message from the wal. It can be called
+// concurrently.
+func (w *BatchWriter) ProcessWALEvent(ctx context.Context, walEvent *wal.Event) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			w.logger.Panic("[PANIC] Panic while processing replication event", loglib.Fields{
+				"wal_data":    walEvent,
+				"panic":       r,
+				"stack_trace": debug.Stack(),
+			})
+
+			err = fmt.Errorf("postgres writer: understanding event: %w:  %v", processor.ErrPanic, r)
+		}
+	}()
+
+	query, err := w.adapter.walEventToQuery(walEvent)
+	if err != nil {
+		return err
+	}
+
+	msg := batch.NewWALMessage(query, walEvent.CommitPosition)
+	return w.batchSender.AddToBatch(ctx, msg)
+}
+
+func (w *BatchWriter) Name() string {
+	return "postgres-batch-writer"
+}
+
+func (w *BatchWriter) Close() error {
+	w.batchSender.Close()
+	return w.pgConn.Close(context.Background())
+}
+
+func (w *BatchWriter) sendBatch(ctx context.Context, batch *batch.Batch[*query]) error {
+	err := w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
+		for _, q := range batch.GetMessages() {
+			if _, err := tx.Exec(ctx, q.sql, q.args...); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	if w.checkpointer != nil {
+		return w.checkpointer(ctx, batch.GetCommitPositions())
+	}
+
+	return nil
+}
diff --git a/pkg/wal/processor/postgres/postgres_query_msg.go b/pkg/wal/processor/postgres/postgres_query_msg.go
new file mode 100644
index 0000000..ff85365
--- /dev/null
+++ b/pkg/wal/processor/postgres/postgres_query_msg.go
@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+type query struct {
+	sql  string
+	args []any
+}
+
+// size returns the size of the message sql query (does not include the
+// parameters)
+func (m *query) Size() int {
+	return len(m.sql)
+}
+
+func (m *query) IsEmpty() bool {
+	return m == nil || m.sql == ""
+}
diff --git a/pkg/wal/processor/postgres/postgres_wal_adapter.go b/pkg/wal/processor/postgres/postgres_wal_adapter.go
new file mode 100644
index 0000000..c31bff9
--- /dev/null
+++ b/pkg/wal/processor/postgres/postgres_wal_adapter.go
@@ -0,0 +1,148 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+
+	"github.com/lib/pq"
+	"github.com/xataio/pgstream/pkg/wal"
+	"github.com/xataio/pgstream/pkg/wal/processor"
+)
+
+type walAdapter interface {
+	walEventToQuery(e *wal.Event) (*query, error)
+}
+
+type adapter struct{}
+
+func (a *adapter) walEventToQuery(e *wal.Event) (*query, error) {
+	// TODO: support DDL parsing
+	if e.Data == nil || processor.IsSchemaLogEvent(e.Data) {
+		return &query{}, nil
+	}
+
+	// DML events
+	switch e.Data.Action {
+	case "I", "U", "D", "T":
+		return a.walDataToQuery(e.Data), nil
+	default:
+		// unsupported actions
+		return &query{}, nil
+	}
+}
+
+func (a *adapter) walDataToQuery(d *wal.Data) *query {
+	switch d.Action {
+	case "T":
+		return a.buildTruncateQuery(d)
+	case "D":
+		return a.buildDeleteQuery(d)
+	case "I":
+		return a.buildInsertQuery(d)
+	case "U":
+		return a.buildUpdateQuery(d)
+	default:
+		return nil
+	}
+}
+
+func (a *adapter) buildTruncateQuery(d *wal.Data) *query {
+	return &query{
+		sql: fmt.Sprintf("TRUNCATE %s", quotedTableName(d.Schema, d.Table)),
+	}
+}
+
+func (a *adapter) buildDeleteQuery(d *wal.Data) *query {
+	primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Identity)
+	whereQuery, whereValues := a.buildWhereQuery(primaryKeyCols, 0)
+	return &query{
+		sql:  fmt.Sprintf("DELETE FROM %s %s", quotedTableName(d.Schema, d.Table), whereQuery),
+		args: whereValues,
+	}
+}
+
+func (a *adapter) buildInsertQuery(d *wal.Data) *query {
+	names := make([]string, 0, len(d.Columns))
+	values := make([]any, 0, len(d.Columns))
+	placeholders := make([]string, 0, len(d.Columns))
+	for i, col := range d.Columns {
+		names = append(names, col.Name)
+		values = append(values, col.Value)
+		placeholders = append(placeholders, fmt.Sprintf("$%d", i+1))
+	}
+
+	return &query{
+		sql: fmt.Sprintf("INSERT INTO %s(%s) VALUES(%s) %s",
+			quotedTableName(d.Schema, d.Table), strings.Join(names, ", "),
+			strings.Join(placeholders, ", "),
+			a.buildOnConflictQuery(d)),
+		args: values,
+	}
+}
+
+func (a *adapter) buildUpdateQuery(d *wal.Data) *query {
+	setQuery, setValues := a.buildSetQuery(d.Columns)
+	primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Columns)
+	whereQuery, whereValues := a.buildWhereQuery(primaryKeyCols, len(d.Columns))
+	return &query{
+		sql:  fmt.Sprintf("UPDATE %s %s %s", quotedTableName(d.Schema, d.Table), setQuery, whereQuery),
+		args: append(setValues, whereValues...),
+	}
+}
+
+func (a *adapter) buildWhereQuery(cols []wal.Column, placeholderOffset int) (string, []any) {
+	whereQuery := "WHERE"
+	whereValues := make([]any, 0, len(cols))
+	for i, c := range cols {
+		if i != 0 {
+			whereQuery = fmt.Sprintf("%s AND", whereQuery)
+		}
+		whereQuery = fmt.Sprintf("%s %s = $%d", whereQuery, c.Name, i+placeholderOffset+1)
+		whereValues = append(whereValues, c.Value)
+	}
+	return whereQuery, whereValues
+}
+
+func (a *adapter) buildSetQuery(cols []wal.Column) (string, []any) {
+	setQuery := "SET"
+	setValues := make([]any, 0, len(cols))
+	for i, c := range cols {
+		if i != 0 {
+			setQuery = fmt.Sprintf("%s,", setQuery)
+		}
+		setQuery = fmt.Sprintf("%s %s = $%d", setQuery, c.Name, i+1)
+		setValues = append(setValues, c.Value)
+	}
+	return setQuery, setValues
+}
+
+func (a *adapter) buildOnConflictQuery(d *wal.Data) string {
+	primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Columns)
+	conflictCols := []string{}
+	for _, col := range primaryKeyCols {
+		conflictCols = append(conflictCols, col.Name)
+	}
+	if len(conflictCols) == 0 {
+		return ""
+	}
+	return fmt.Sprintf("ON CONFLICT (%s) DO NOTHING", strings.Join(conflictCols, ", "))
+}
+
+func (a *adapter) extractPrimaryKeyColumns(colIDs []string, cols []wal.Column) []wal.Column {
+	primaryKeyColumns := make([]wal.Column, 0, len(colIDs))
+	for _, col := range cols {
+		if !slices.Contains(colIDs, col.ID) {
+			continue
+		}
+		primaryKeyColumns = append(primaryKeyColumns, col)
+	}
+
+	return primaryKeyColumns
+}
+
+func quotedTableName(schemaName, tableName string) string {
+	return fmt.Sprintf("%s.%s", pq.QuoteIdentifier(schemaName), pq.QuoteIdentifier(tableName))
+}


Test patch:
diff --git a/pkg/wal/processor/postgres/helper_test.go b/pkg/wal/processor/postgres/helper_test.go
new file mode 100644
index 0000000..12332d6
--- /dev/null
+++ b/pkg/wal/processor/postgres/helper_test.go
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "github.com/xataio/pgstream/pkg/wal"
+
+type mockAdapter struct {
+	walEventToQueryFn func(*wal.Event) (*query, error)
+}
+
+func (m *mockAdapter) walEventToQuery(e *wal.Event) (*query, error) {
+	return m.walEventToQueryFn(e)
+}
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer_test.go b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
new file mode 100644
index 0000000..d22e106
--- /dev/null
+++ b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
@@ -0,0 +1,310 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/wal"
+	"github.com/xataio/pgstream/pkg/wal/checkpointer"
+	"github.com/xataio/pgstream/pkg/wal/processor"
+	"github.com/xataio/pgstream/pkg/wal/processor/batch"
+	batchmocks "github.com/xataio/pgstream/pkg/wal/processor/batch/mocks"
+)
+
+var (
+	testSchema = "test_schema"
+	testTable  = "test_table"
+
+	testLSNStr         = "1/CF54A048"
+	testCommitPosition = wal.CommitPosition(testLSNStr)
+
+	errTest = errors.New("oh noes")
+)
+
+func TestBatchWriter_ProcessWALEvent(t *testing.T) {
+	t.Parallel()
+
+	testWalEvent := &wal.Event{
+		Data: &wal.Data{
+			Action: "I",
+			LSN:    testLSNStr,
+			Schema: testSchema,
+			Table:  testTable,
+		},
+		CommitPosition: testCommitPosition,
+	}
+
+	testInsertQuery := &query{
+		sql:  "INSERT INTO test(id, name) VALUES($1, $2) ON CONFLICT (id) DO NOTHING",
+		args: []any{1, "alice"},
+	}
+
+	tests := []struct {
+		name        string
+		walEvent    *wal.Event
+		batchSender *batchmocks.BatchSender[*query]
+		adapter     walAdapter
+
+		wantMsgs []*batch.WALMessage[*query]
+		wantErr  error
+	}{
+		{
+			name:        "ok",
+			walEvent:    testWalEvent,
+			batchSender: batchmocks.NewBatchSender[*query](),
+			adapter: &mockAdapter{
+				walEventToQueryFn: func(e *wal.Event) (*query, error) {
+					require.Equal(t, e, testWalEvent)
+					return testInsertQuery, nil
+				},
+			},
+
+			wantMsgs: []*batch.WALMessage[*query]{
+				batch.NewWALMessage(testInsertQuery, testCommitPosition),
+			},
+			wantErr: nil,
+		},
+		{
+			name:        "error - event to query",
+			walEvent:    testWalEvent,
+			batchSender: batchmocks.NewBatchSender[*query](),
+			adapter: &mockAdapter{
+				walEventToQueryFn: func(e *wal.Event) (*query, error) {
+					return nil, errTest
+				},
+			},
+
+			wantMsgs: []*batch.WALMessage[*query]{},
+			wantErr:  errTest,
+		},
+		{
+			name:     "error - adding to batch",
+			walEvent: testWalEvent,
+			batchSender: func() *batchmocks.BatchSender[*query] {
+				s := batchmocks.NewBatchSender[*query]()
+				s.AddToBatchFn = func(ctx context.Context, w *batch.WALMessage[*query]) error { return errTest }
+				return s
+			}(),
+			adapter: &mockAdapter{
+				walEventToQueryFn: func(e *wal.Event) (*query, error) {
+					require.Equal(t, e, testWalEvent)
+					return testInsertQuery, nil
+				},
+			},
+
+			wantMsgs: []*batch.WALMessage[*query]{},
+			wantErr:  errTest,
+		},
+		{
+			name:        "error - panic recovery",
+			walEvent:    testWalEvent,
+			batchSender: batchmocks.NewBatchSender[*query](),
+			adapter: &mockAdapter{
+				walEventToQueryFn: func(e *wal.Event) (*query, error) {
+					panic(errTest)
+				},
+			},
+
+			wantMsgs: []*batch.WALMessage[*query]{},
+			wantErr:  processor.ErrPanic,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			writer := &BatchWriter{
+				logger:      loglib.NewNoopLogger(),
+				batchSender: tc.batchSender,
+				adapter:     tc.adapter,
+			}
+
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			go func() {
+				defer tc.batchSender.Close()
+				err := writer.ProcessWALEvent(ctx, tc.walEvent)
+				if !errors.Is(err, tc.wantErr) {
+					require.Equal(t, err.Error(), tc.wantErr.Error())
+				}
+			}()
+
+			msgs := tc.batchSender.GetWALMessages()
+			require.Equal(t, tc.wantMsgs, msgs)
+		})
+	}
+}
+
+func TestBatchWriter_sendBatch(t *testing.T) {
+	t.Parallel()
+
+	testQuery := &query{
+		sql:  "INSERT INTO test(id, name) VALUES($1, $2)",
+		args: []any{1, "alice"},
+	}
+
+	tests := []struct {
+		name         string
+		pgconn       *pgmocks.Querier
+		batchSender  *batchmocks.BatchSender[*query]
+		batch        *batch.Batch[*query]
+		checkpointer checkpointer.Checkpoint
+
+		wantErr error
+	}{
+		{
+			name: "ok",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							require.Equal(t, testQuery.sql, query)
+							require.Equal(t, testQuery.args, args)
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+				CloseFn: func(ctx context.Context) error { return nil },
+			},
+			batchSender: batchmocks.NewBatchSender[*query](),
+			batch:       batch.NewBatch([]*query{testQuery}, []wal.CommitPosition{testCommitPosition}),
+
+			wantErr: nil,
+		},
+		{
+			name: "error - executing query",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							return pglib.CommandTag{}, errTest
+						},
+					}
+					return f(&mockTx)
+				},
+				CloseFn: func(ctx context.Context) error { return nil },
+			},
+			batchSender: batchmocks.NewBatchSender[*query](),
+			batch:       batch.NewBatch([]*query{testQuery}, []wal.CommitPosition{testCommitPosition}),
+
+			wantErr: errTest,
+		},
+		{
+			name: "error - checkpointing",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+				CloseFn: func(ctx context.Context) error { return nil },
+			},
+			batchSender:  batchmocks.NewBatchSender[*query](),
+			batch:        batch.NewBatch([]*query{testQuery}, []wal.CommitPosition{testCommitPosition}),
+			checkpointer: func(ctx context.Context, positions []wal.CommitPosition) error { return errTest },
+
+			wantErr: errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			writer := &BatchWriter{
+				logger:       loglib.NewNoopLogger(),
+				batchSender:  tc.batchSender,
+				pgConn:       tc.pgconn,
+				checkpointer: tc.checkpointer,
+			}
+			defer writer.Close()
+
+			err := writer.sendBatch(context.Background(), tc.batch)
+			require.ErrorIs(t, err, tc.wantErr)
+		})
+	}
+}
+
+func TestBatchWriter(t *testing.T) {
+	t.Parallel()
+
+	conn := &pgmocks.Querier{
+		ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+			mockTx := pgmocks.Tx{
+				ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+					time.Sleep(time.Second)
+					return pglib.CommandTag{}, errTest
+				},
+			}
+			return f(&mockTx)
+		},
+		CloseFn: func(ctx context.Context) error { return nil },
+	}
+
+	adapter := &mockAdapter{
+		walEventToQueryFn: func(e *wal.Event) (*query, error) {
+			return &query{
+				sql:  "INSERT INTO test(id, name) VALUES($1, $2) ON CONFLICT (id) DO NOTHING",
+				args: []any{1, "alice"},
+			}, nil
+		},
+	}
+
+	bw := BatchWriter{
+		logger:  loglib.NewNoopLogger(),
+		pgConn:  conn,
+		adapter: adapter,
+	}
+	defer bw.Close()
+
+	var err error
+	bw.batchSender, err = batch.NewSender(&batch.Config{
+		BatchTimeout:  time.Second,
+		MaxBatchSize:  10,
+		MaxBatchBytes: 10000,
+	}, bw.sendBatch, bw.logger)
+	require.NoError(t, err)
+
+	doneChan := make(chan struct{}, 1)
+	go func() {
+		err := bw.batchSender.Send(context.Background())
+		require.ErrorIs(t, err, errTest)
+		doneChan <- struct{}{}
+		close(doneChan)
+	}()
+
+	timer := time.NewTimer(5 * time.Second)
+	defer timer.Stop()
+	var processErr error
+	for {
+		select {
+		case <-doneChan:
+			require.ErrorIs(t, processErr, errTest)
+			return
+		case <-timer.C:
+			t.Error("test timeout")
+			return
+		default:
+			processErr = bw.ProcessWALEvent(context.Background(), &wal.Event{
+				CommitPosition: wal.CommitPosition("1"),
+				Data: &wal.Data{
+					Action: "I",
+				},
+			})
+		}
+	}
+}
diff --git a/pkg/wal/processor/postgres/postgres_wal_adapter_test.go b/pkg/wal/processor/postgres/postgres_wal_adapter_test.go
new file mode 100644
index 0000000..4a3289d
--- /dev/null
+++ b/pkg/wal/processor/postgres/postgres_wal_adapter_test.go
@@ -0,0 +1,237 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/rs/xid"
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/pkg/schemalog"
+	"github.com/xataio/pgstream/pkg/wal"
+)
+
+func TestAdapter_walEventToQuery(t *testing.T) {
+	t.Parallel()
+
+	testCommitPosition := wal.CommitPosition("1/0")
+	testTable := "table"
+	testSchema := "test"
+
+	tests := []struct {
+		name  string
+		event *wal.Event
+
+		wantQuery *query
+		wantErr   error
+	}{
+		{
+			name: "ok - no data",
+			event: &wal.Event{
+				CommitPosition: testCommitPosition,
+			},
+
+			wantQuery: &query{},
+			wantErr:   nil,
+		},
+		{
+			name: "ok - schema event",
+			event: &wal.Event{
+				Data: &wal.Data{
+					Schema: schemalog.SchemaName,
+					Table:  schemalog.TableName,
+				},
+				CommitPosition: testCommitPosition,
+			},
+
+			wantQuery: &query{},
+			wantErr:   nil,
+		},
+		{
+			name: "ok - supported data event",
+			event: &wal.Event{
+				Data: &wal.Data{
+					Action: "T",
+					Schema: testSchema,
+					Table:  testTable,
+				},
+				CommitPosition: testCommitPosition,
+			},
+
+			wantQuery: &query{
+				sql: fmt.Sprintf("TRUNCATE %s", quotedTableName(testSchema, testTable)),
+			},
+			wantErr: nil,
+		},
+		{
+			name: "ok - unsupported data event",
+			event: &wal.Event{
+				Data: &wal.Data{
+					Action: "X",
+					Schema: testSchema,
+					Table:  testTable,
+				},
+				CommitPosition: testCommitPosition,
+			},
+
+			wantQuery: &query{},
+			wantErr:   nil,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			a := &adapter{}
+			query, err := a.walEventToQuery(tc.event)
+			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantQuery, query)
+		})
+	}
+}
+
+func TestAdapter_walDataToQuery(t *testing.T) {
+	t.Parallel()
+
+	testTableID := xid.New()
+	testTable := "table"
+	testSchema := "test"
+	quotedTestTable := quotedTableName(testSchema, testTable)
+
+	columnID := func(i int) string {
+		return fmt.Sprintf("%s-%d", testTableID, i)
+	}
+
+	tests := []struct {
+		name    string
+		walData *wal.Data
+
+		wantQuery *query
+	}{
+		{
+			name: "truncate",
+			walData: &wal.Data{
+				Action: "T",
+				Schema: testSchema,
+				Table:  testTable,
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1)},
+				},
+			},
+
+			wantQuery: &query{
+				sql: fmt.Sprintf("TRUNCATE %s", quotedTestTable),
+			},
+		},
+		{
+			name: "delete with simple primary key",
+			walData: &wal.Data{
+				Action: "D",
+				Schema: testSchema,
+				Table:  testTable,
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1)},
+				},
+			},
+
+			wantQuery: &query{
+				sql:  fmt.Sprintf("DELETE FROM %s WHERE id = $1", quotedTestTable),
+				args: []any{1},
+			},
+		},
+		{
+			name: "delete with composite primary key",
+			walData: &wal.Data{
+				Action: "D",
+				Schema: testSchema,
+				Table:  testTable,
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1), columnID(2)},
+				},
+			},
+
+			wantQuery: &query{
+				sql:  fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND name = $2", quotedTestTable),
+				args: []any{1, "alice"},
+			},
+		},
+		{
+			name: "insert",
+			walData: &wal.Data{
+				Action: "I",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1)},
+				},
+			},
+
+			wantQuery: &query{
+				sql:  fmt.Sprintf("INSERT INTO %s(id, name) VALUES($1, $2) ON CONFLICT (id) DO NOTHING", quotedTestTable),
+				args: []any{1, "alice"},
+			},
+		},
+		{
+			name: "update",
+			walData: &wal.Data{
+				Action: "U",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1)},
+				},
+			},
+
+			wantQuery: &query{
+				sql:  fmt.Sprintf("UPDATE %s SET id = $1, name = $2 WHERE id = $3", quotedTestTable),
+				args: []any{1, "alice", 1},
+			},
+		},
+		{
+			name: "unknown",
+			walData: &wal.Data{
+				Action: "X",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{
+					InternalColIDs: []string{columnID(1)},
+				},
+			},
+
+			wantQuery: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			a := &adapter{}
+			query := a.walDataToQuery(tc.walData)
+			require.Equal(t, tc.wantQuery, query)
+		})
+	}
+}