Repository: xataio/pgstream
Pull request #417 (2025-06-24 11:52:10)

[Postgres Writer] Create schema if not exists
Currently any non `public` schemas are not replicated to postgres, since the `CREATE SCHEMA` query is not executed. This affects both the schema snapshot in `schemalog` mode and the replication to postgres.

This PR updates the DDL adapter to create the schema if it doesn't exist when receiving a DDL event. It also fixes the condition to not add `DEFAULT` statements in table creation to generated columns.


Patch:
diff --git a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
index 2ae6086..af5ea42 100644
--- a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
+++ b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 
 	"github.com/xataio/pgstream/internal/json"
+	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/snapshot"
 	"github.com/xataio/pgstream/pkg/snapshot/generator"
@@ -17,6 +18,7 @@ type SnapshotGenerator struct {
 	marshaler      func(any) ([]byte, error)
 	processRow     snapshot.RowProcessor
 	generator      generator.SnapshotGenerator
+	logger         loglib.Logger
 }
 
 type Option func(*SnapshotGenerator)
@@ -26,6 +28,7 @@ func NewSnapshotGenerator(schemalogStore schemalog.Store, processRow snapshot.Ro
 		schemalogStore: schemalogStore,
 		processRow:     processRow,
 		marshaler:      json.Marshal,
+		logger:         loglib.NewNoopLogger(),
 	}
 
 	for _, opt := range opts {
@@ -41,15 +44,32 @@ func WithSnapshotGenerator(g generator.SnapshotGenerator) Option {
 	}
 }
 
+func WithLogger(logger loglib.Logger) Option {
+	return func(sg *SnapshotGenerator) {
+		sg.logger = loglib.NewLogger(logger).WithFields(loglib.Fields{
+			loglib.ModuleField: "postgres_schemalog_snapshot_generator",
+		})
+	}
+}
+
 func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
 	snapshotErrs := make(snapshot.Errors)
 	for schema := range ss.SchemaTables {
+		s.logger.Info("creating schema snapshot", loglib.Fields{"schema": schema, "tables": ss.SchemaTables[schema]})
 		err := func() error {
 			logEntry, err := s.schemalogStore.Insert(ctx, schema)
 			if err != nil {
 				return err
 			}
 
+			s.logger.Debug("created schema log entry", loglib.Fields{
+				"schema":     schema,
+				"log_entry":  logEntry.ID,
+				"version":    logEntry.Version,
+				"created_at": logEntry.CreatedAt,
+				"acked":      logEntry.Acked,
+			})
+
 			row, err := s.logEntryToSnapshotRow(logEntry)
 			if err != nil {
 				return err
diff --git a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
index 6850157..2a7b43f 100644
--- a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
+++ b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
@@ -124,7 +124,8 @@ func newSchemaSnapshotGenerator(ctx context.Context, cfg *SchemaSnapshotConfig,
 		return schemalogsnapshotgenerator.NewSnapshotGenerator(
 			schemaLogStore,
 			processRow,
-			schemalogsnapshotgenerator.WithSnapshotGenerator(g)), nil
+			schemalogsnapshotgenerator.WithSnapshotGenerator(g),
+			schemalogsnapshotgenerator.WithLogger(logger)), nil
 	case cfg.DumpRestore != nil:
 		// postgres pgdump/pgrestore schema snapshot generator
 		opts := []pgdumprestoregenerator.Option{
diff --git a/pkg/wal/processor/postgres/postgres_wal_ddl_adapter.go b/pkg/wal/processor/postgres/postgres_wal_ddl_adapter.go
index d3549fe..141485c 100644
--- a/pkg/wal/processor/postgres/postgres_wal_ddl_adapter.go
+++ b/pkg/wal/processor/postgres/postgres_wal_ddl_adapter.go
@@ -46,7 +46,23 @@ func (a *ddlAdapter) schemaLogToQueries(ctx context.Context, schemaLog *schemalo
 
 	diff := a.schemaDiffer(previousSchemaLog, schemaLog)
 
-	return a.schemaDiffToQueries(schemaLog.SchemaName, diff)
+	queries := []*query{
+		a.createSchemaIfNotExists(schemaLog.SchemaName),
+	}
+
+	schemaQueries, err := a.schemaDiffToQueries(schemaLog.SchemaName, diff)
+	if err != nil {
+		return nil, err
+	}
+
+	return append(queries, schemaQueries...), nil
+}
+
+const createSchemaIfNotExistsQuery = "CREATE SCHEMA IF NOT EXISTS %s"
+
+func (a *ddlAdapter) createSchemaIfNotExists(schemaName string) *query {
+	createSchemaQuery := fmt.Sprintf(createSchemaIfNotExistsQuery, pglib.QuoteIdentifier(schemaName))
+	return a.newDDLQuery(schemaName, "", createSchemaQuery)
 }
 
 func (a *ddlAdapter) schemaDiffToQueries(schemaName string, diff *schemalog.Diff) ([]*query, error) {
@@ -108,9 +124,10 @@ func (a *ddlAdapter) buildColumnDefinition(column *schemalog.Column) string {
 	if !column.Nullable {
 		colDefinition = fmt.Sprintf("%s NOT NULL", colDefinition)
 	}
-	// do not set default values with sequences since they will differ between
-	// source/target. Keep source database as source of truth.
-	if column.DefaultValue != nil && (!strings.Contains(*column.DefaultValue, "seq") || column.Generated) {
+	// do not set default values with sequences and generated columns since they
+	// must be aligned between source/target. Keep source database as source of
+	// truth.
+	if column.DefaultValue != nil && !strings.Contains(*column.DefaultValue, "seq") && !column.Generated {
 		colDefinition = fmt.Sprintf("%s DEFAULT %s", colDefinition, *column.DefaultValue)
 	}
 


Test patch:
diff --git a/pkg/wal/processor/postgres/postgres_wal_ddl_adapter_test.go b/pkg/wal/processor/postgres/postgres_wal_ddl_adapter_test.go
index c9798f2..84f800c 100644
--- a/pkg/wal/processor/postgres/postgres_wal_ddl_adapter_test.go
+++ b/pkg/wal/processor/postgres/postgres_wal_ddl_adapter_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/rs/xid"
 	"github.com/stretchr/testify/require"
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	schemalogmocks "github.com/xataio/pgstream/pkg/schemalog/mocks"
 )
@@ -65,6 +66,11 @@ func TestDDLAdapter_walDataToQueries(t *testing.T) {
 			logEntry: testLogEntry(0),
 
 			wantQueries: []*query{
+				{
+					schema: testSchema,
+					sql:    fmt.Sprintf(createSchemaIfNotExistsQuery, pglib.QuoteIdentifier(testSchema)),
+					isDDL:  true,
+				},
 				{
 					schema: testSchema,
 					table:  table2,
@@ -99,6 +105,11 @@ func TestDDLAdapter_walDataToQueries(t *testing.T) {
 			logEntry: testLogEntry(1),
 
 			wantQueries: []*query{
+				{
+					schema: testSchema,
+					sql:    fmt.Sprintf(createSchemaIfNotExistsQuery, pglib.QuoteIdentifier(testSchema)),
+					isDDL:  true,
+				},
 				{
 					schema: testSchema,
 					table:  table2,