Repository: xataio/pgstream
Pull request #103 (2024-12-10 15:06:18)

Add snapshot aggregator
This PR adds a snapshot generator aggregator. This will allow us to combine both schema and data snapshotting when configured, as well as allowing other snapshot generator implementations to be added on later on.


Patch:
diff --git a/pkg/snapshot/generator/snapshot_generator_aggregator.go b/pkg/snapshot/generator/snapshot_generator_aggregator.go
new file mode 100644
index 0000000..c0b55f9
--- /dev/null
+++ b/pkg/snapshot/generator/snapshot_generator_aggregator.go
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package generator
+
+import (
+	"context"
+	"errors"
+
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+// GeneratorAggregator aggregates snapshot generators
+type GeneratorAggregator struct {
+	generators []SnapshotGenerator
+}
+
+// NewAggregator will aggregate the generators on input.
+func NewAggregator(generators []SnapshotGenerator) *GeneratorAggregator {
+	return &GeneratorAggregator{
+		generators: generators,
+	}
+}
+
+// CreateSnapshot will call the CreateSnapshot method for each of the aggregated
+// snapshot generators sequentially, in the order in which they were provided.
+// If one of them fails, an error will be returned without continuing with the
+// rest of generator calls.
+func (a *GeneratorAggregator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+	for _, gen := range a.generators {
+		if err := gen.CreateSnapshot(ctx, ss); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Close will aggregate the errors of the generator close calls.
+func (a *GeneratorAggregator) Close() error {
+	var closeErrs error
+	for _, gen := range a.generators {
+		if err := gen.Close(); err != nil {
+			if closeErrs != nil {
+				closeErrs = errors.Join(closeErrs, err)
+			} else {
+				closeErrs = err
+			}
+		}
+	}
+	return closeErrs
+}


Test patch:
diff --git a/pkg/snapshot/generator/snapshot_generator_aggregator_test.go b/pkg/snapshot/generator/snapshot_generator_aggregator_test.go
new file mode 100644
index 0000000..69b6420
--- /dev/null
+++ b/pkg/snapshot/generator/snapshot_generator_aggregator_test.go
@@ -0,0 +1,105 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package generator
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+func TestAggregator_CreateSnapshot(t *testing.T) {
+	t.Parallel()
+
+	testSnapshot := &snapshot.Snapshot{
+		SchemaName: "test-schema",
+		TableNames: []string{"table1", "table2"},
+	}
+
+	newGenerator := func(err error) *mockGenerator {
+		return &mockGenerator{
+			createSnapshotFn: func(ctx context.Context, ss *snapshot.Snapshot) error {
+				require.Equal(t, testSnapshot, ss)
+				return err
+			},
+		}
+	}
+
+	errTest := errors.New("oh noes")
+
+	tests := []struct {
+		name       string
+		generators []SnapshotGenerator
+
+		wantErr error
+	}{
+		{
+			name:       "ok",
+			generators: []SnapshotGenerator{newGenerator(nil), newGenerator(nil)},
+			wantErr:    nil,
+		},
+		{
+			name:       "error on first generator",
+			generators: []SnapshotGenerator{newGenerator(errTest), newGenerator(errors.New("should not be called"))},
+			wantErr:    errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			agg := NewAggregator(tc.generators)
+
+			err := agg.CreateSnapshot(context.Background(), testSnapshot)
+			require.Equal(t, err, tc.wantErr)
+		})
+	}
+}
+
+func TestAggregator_Close(t *testing.T) {
+	t.Parallel()
+
+	newGenerator := func(err error) *mockGenerator {
+		return &mockGenerator{
+			closeFn: func() error {
+				return err
+			},
+		}
+	}
+
+	errTest1 := errors.New("oh noes")
+	errTest2 := errors.New("oh nose")
+
+	tests := []struct {
+		name       string
+		generators []SnapshotGenerator
+
+		wantErr error
+	}{
+		{
+			name:       "ok",
+			generators: []SnapshotGenerator{newGenerator(nil), newGenerator(nil)},
+			wantErr:    nil,
+		},
+		{
+			name:       "error on first generator",
+			generators: []SnapshotGenerator{newGenerator(errTest1), newGenerator(errTest2)},
+			wantErr:    errors.Join(errTest1, errTest2),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			agg := NewAggregator(tc.generators)
+
+			err := agg.Close()
+			require.Equal(t, err, tc.wantErr)
+		})
+	}
+}