Clean up code.

sources/dynamodb/snapshot.go

@@ -71,6 +71,7 @@ func (s *SnapshotStore) streamAndPublish(ctx context.Context, writer writers.Wri
 		return fmt.Errorf("failed to retrieve primary keys: %w", err)
 	}
 
+	writer.SetRunOnComplete(false)
 	for _, file := range s.cfg.SnapshotSettings.SpecifiedFiles {
 		logFields := []any{
 			slog.String("fileName", *file.Key),
@@ -101,5 +102,5 @@ func (s *SnapshotStore) streamAndPublish(ctx context.Context, writer writers.Wri
 		slog.Info("Successfully processed file...", logFields...)
 	}
 
-	return nil
+	return writer.OnComplete()
 }

writers/writer.go

@@ -18,10 +18,15 @@ type DestinationWriter interface {
 type Writer struct {
 	destinationWriter DestinationWriter
 	logProgress       bool
+	runOnComplete     bool
+}
+
+func (w *Writer) SetRunOnComplete(runOnComplete bool) {
+	w.runOnComplete = runOnComplete
 }
 
 func New(destinationWriter DestinationWriter, logProgress bool) Writer {
-	return Writer{destinationWriter, logProgress}
+	return Writer{destinationWriter: destinationWriter, logProgress: logProgress, runOnComplete: true}
 }
 
 // Write writes all the messages from an iterator to the destination.
@@ -55,9 +60,19 @@ func (w *Writer) Write(ctx context.Context, iter iterator.Iterator[[]lib.RawMess
 		}
 	}
 
-	if err := w.destinationWriter.OnComplete(); err != nil {
-		return 0, fmt.Errorf("failed running destination OnComplete: %w", err)
+	if w.runOnComplete {
+		if err := w.destinationWriter.OnComplete(); err != nil {
+			return 0, fmt.Errorf("failed running destination OnComplete: %w", err)
+		}
 	}
 
 	return count, nil
 }
+
+func (w *Writer) OnComplete() error {
+	if err := w.destinationWriter.OnComplete(); err != nil {
+		return fmt.Errorf("failed running destination OnComplete: %w", err)
+	}
+
+	return nil
+}