snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

message/decorator.go

@@ -0,0 +1,81 @@
+package message
+
+import (
+	"context"
+	"sync"
+)
+
+// MessageTransformSubscriberDecorator creates a subscriber decorator that calls transform
+// on each message that passes through the subscriber.
+func MessageTransformSubscriberDecorator(transform func(*Message)) SubscriberDecorator {
+	if transform == nil {
+		panic("transform function is nil")
+	}
+	return func(sub Subscriber) (Subscriber, error) {
+		return &messageTransformSubscriberDecorator{
+			sub:       sub,
+			transform: transform,
+		}, nil
+	}
+}
+
+// MessageTransformPublisherDecorator creates a publisher decorator that calls transform
+// on each message that passes through the publisher.
+func MessageTransformPublisherDecorator(transform func(*Message)) PublisherDecorator {
+	if transform == nil {
+		panic("transform function is nil")
+	}
+	return func(pub Publisher) (Publisher, error) {
+		return &messageTransformPublisherDecorator{
+			Publisher: pub,
+			transform: transform,
+		}, nil
+	}
+}
+
+type messageTransformSubscriberDecorator struct {
+	sub Subscriber
+
+	transform   func(*Message)
+	subscribeWg sync.WaitGroup
+}
+
+func (t *messageTransformSubscriberDecorator) Subscribe(ctx context.Context, topic string) (<-chan *Message, error) {
+	in, err := t.sub.Subscribe(ctx, topic)
+	if err != nil {
+		return nil, err
+	}
+
+	out := make(chan *Message)
+	t.subscribeWg.Add(1)
+	go func() {
+		for msg := range in {
+			t.transform(msg)
+			out <- msg
+		}
+		close(out)
+		t.subscribeWg.Done()
+	}()
+
+	return out, nil
+}
+
+func (t *messageTransformSubscriberDecorator) Close() error {
+	err := t.sub.Close()
+
+	t.subscribeWg.Wait()
+	return err
+}
+
+type messageTransformPublisherDecorator struct {
+	Publisher
+	transform func(*Message)
+}
+
+// Publish applies the transform to each message and returns the underlying Publisher's result.
+func (d messageTransformPublisherDecorator) Publish(topic string, messages ...*Message) error {
+	for i := range messages {
+		d.transform(messages[i])
+	}
+	return d.Publisher.Publish(topic, messages...)
+}