57 lines
1.4 KiB
Go
57 lines
1.4 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
|
|
"github.com/ThreeDotsLabs/watermill/components/cqrs"
|
|
"github.com/ThreeDotsLabs/watermill/message"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
func GenerateMessageMetadata(partitionKey string) *MessageMetadata {
|
|
return &MessageMetadata{
|
|
PartitionKey: partitionKey,
|
|
CreatedAt: timestamppb.Now(),
|
|
}
|
|
}
|
|
|
|
type CqrsMarshalerDecorator struct {
|
|
cqrs.ProtoMarshaler
|
|
}
|
|
|
|
const PartitionKeyMetadataField = "partition_key"
|
|
|
|
func (c CqrsMarshalerDecorator) Marshal(v interface{}) (*message.Message, error) {
|
|
msg, err := c.ProtoMarshaler.Marshal(v)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
pm, ok := v.(ProtoMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("%T does not implement ProtoMessage and can't be marshaled", v)
|
|
}
|
|
|
|
metadata := pm.GetMetadata()
|
|
if metadata == nil {
|
|
return nil, fmt.Errorf("%T.GetMetadata returned nil", v)
|
|
}
|
|
|
|
msg.Metadata.Set(PartitionKeyMetadataField, metadata.PartitionKey)
|
|
msg.Metadata.Set("created_at", metadata.CreatedAt.AsTime().String())
|
|
|
|
return msg, nil
|
|
}
|
|
|
|
type ProtoMessage interface {
|
|
GetMetadata() *MessageMetadata
|
|
}
|
|
|
|
// GenerateKafkaPartitionKey is a function that generates a partition key for Kafka messages.
|
|
func GenerateKafkaPartitionKey(topic string, msg *message.Message) (string, error) {
|
|
slog.Debug("Setting partition key", "topic", topic, "msg_metadata", msg.Metadata)
|
|
|
|
return msg.Metadata.Get(PartitionKeyMetadataField), nil
|
|
}
|