feat: Add support to set partition key and headers via Request Context

src/Paramore.Brighter.Transformers.JustSaying/JustSayingMessageMapper.cs

@@ -1,11 +1,14 @@
 using System;
+using System.Collections.Generic;
 using System.Net.Mime;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Paramore.Brighter.Extensions;
 using Paramore.Brighter.JsonConverters;
 using Paramore.Brighter.Transformers.JustSaying.Extensions;
 using Paramore.Brighter.Transformers.JustSaying.JsonConverters;
+using Paramore.Brighter.Transforms.Attributes;
 
 namespace Paramore.Brighter.Transformers.JustSaying;
 
@@ -38,6 +41,7 @@ public JustSayingMessageMapper()
     public IRequestContext? Context { get; set; }
 
     /// <inheritdoc />
+    [CloudEvents(0)]
     public Task<Message> MapToMessageAsync(TMessage request, Publication publication, CancellationToken cancellationToken = default)
     {
         return Task.FromResult(MapToMessage(request, publication));
@@ -50,6 +54,7 @@ public Task<TMessage> MapToRequestAsync(Message message, CancellationToken cance
     }
 
     /// <inheritdoc />
+    [CloudEvents(0)]
     public Message MapToMessage(TMessage request, Publication publication)
     {
         var messageType = request switch
@@ -64,6 +69,7 @@ public Message MapToMessage(TMessage request, Publication publication)
 
     private Message JustSayingToMessage(TMessage request, MessageType messageType, Publication publication)
     {
+        var defaultHeaders = publication.DefaultHeaders ?? new Dictionary<string, object>();
         var justSaying = (IJustSayingRequest)request;
         justSaying.Id = GetId(justSaying.Id);
         justSaying.Conversation = GetCorrelationId(justSaying.Conversation);
@@ -80,16 +86,19 @@ private Message JustSayingToMessage(TMessage request, MessageType messageType, P
                 messageType: messageType,
                 subject: GetSubject(publication),
                 timeStamp: justSaying.TimeStamp,
-                topic: publication.Topic!),
+                topic: publication.Topic!,
+                partitionKey: Context.GetPartitionKey())
+            {
+                Bag = defaultHeaders.Merge(Context.GetHeaders())
+            },
             new MessageBody(
                 JsonSerializer.SerializeToUtf8Bytes(request, JsonSerialisationOptions.Options),
                 s_justSaying));
-
-
     }
 
     private Message GenericToMessage(TMessage request, MessageType messageType, Publication publication)
     {
+        var defaultHeaders = publication.DefaultHeaders ?? new Dictionary<string, object>();
         var doc = JsonSerializer.SerializeToNode(request, JsonSerialisationOptions.Options)!;
         var messageId = GetId(doc.GetId(nameof(IJustSayingRequest.Id)));
         var correlationId = GetCorrelationId(doc.GetId(nameof(IJustSayingRequest.Conversation))); 
@@ -110,7 +119,11 @@ private Message GenericToMessage(TMessage request, MessageType messageType, Publ
                 messageType: messageType,
                 subject: GetSubject(publication),
                 timeStamp: timestamp,
-                topic: publication.Topic!),
+                topic: publication.Topic!,
+                partitionKey: Context.GetPartitionKey())
+            {
+                Bag = defaultHeaders.Merge(Context.GetHeaders())
+            },
             new MessageBody(
                 doc.ToJsonString(JsonSerialisationOptions.Options),
                 s_justSaying));

src/Paramore.Brighter.Transformers.MassTransit/MassTransitMessageMapper.cs

@@ -7,7 +7,9 @@
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Paramore.Brighter.Extensions;
 using Paramore.Brighter.JsonConverters;
+using Paramore.Brighter.Transforms.Attributes;
 
 namespace Paramore.Brighter.Transformers.MassTransit;
 
@@ -50,6 +52,7 @@ public class MassTransitMessageMapper<TMessage> : IAmAMessageMapper<TMessage>, I
     public IRequestContext? Context { get; set; }
 
     /// <inheritdoc />
+    [CloudEvents(0)]
     public Task<Message> MapToMessageAsync(TMessage request, Publication publication,
         CancellationToken cancellationToken = default)
     {
@@ -63,29 +66,22 @@ public Task<TMessage> MapToRequestAsync(Message message, CancellationToken cance
     }
 
     /// <inheritdoc />
+    [CloudEvents(0)]
     public virtual Message MapToMessage(TMessage request, Publication publication)
     {
         var timestamp = DateTimeOffset.UtcNow;
-        var bag = new Dictionary<string, object>();
-        if (Context is { Bag.Count: > 0 })
-        {
-            foreach (var pair in Context.Bag)
-            {
-                if (!pair.Key.StartsWith(MassTransitHeaderNames.HeaderPrefix))
-                {
-                    bag[pair.Key] = pair.Value;
-                }
-            }
-        }
-
+
+        var defaultHeaders = publication.DefaultHeaders ?? new Dictionary<string, object>();
+        var headers = defaultHeaders.Merge(Context.GetHeaders());
+
         var envelop = new MassTransitMessageEnvelop<TMessage>
         {
             ConversationId = GetConversationId(),
             CorrelationId = GetCorrelationId(),
             DestinationAddress = GetDestinationAddress(),
             ExpirationTime = GetExpirationTime(),
             FaultAddress = GetFaultAddress(),
-            Headers = null,
+            Headers = headers!,
             Host = s_hostInfo,
             InitiatorId = GetInitiatorId(),
             Message = request,
@@ -109,9 +105,10 @@ public virtual Message MapToMessage(TMessage request, Publication publication)
                     _ => MessageType.MT_DOCUMENT
                 },
                 timeStamp: timestamp,
-                topic: publication.Topic!)
+                topic: publication.Topic!,
+                partitionKey: Context.GetPartitionKey())
             {
-                Bag = bag
+                Bag = headers 
             },
             new MessageBody(JsonSerializer.SerializeToUtf8Bytes(envelop, JsonSerialisationOptions.Options),
                 MassTransitContentType)

src/Paramore.Brighter/Extensions/DictionaryExtensions.cs

@@ -0,0 +1,67 @@
+using System.Collections.Generic;
+
+namespace Paramore.Brighter.Extensions;
+
+/// <summary>
+/// Provides extension methods for dictionaries.
+/// </summary>
+public static class DictionaryExtensions
+{
+    /// <summary>
+    /// Merges two dictionaries into a new dictionary, with values from the second dictionary overwriting those from the first.
+    /// </summary>
+    /// <param name="dict1">The primary dictionary to merge (will be copied first)</param>
+    /// <param name="dict2">The secondary dictionary to merge (may be null)</param>
+    /// <returns>
+    /// A new dictionary containing the combined key-value pairs from both dictionaries.
+    /// When keys exist in both dictionaries, the value from <paramref name="dict2"/> takes precedence.
+    /// </returns>
+    /// <remarks>
+    /// <para>
+    /// This method performs a shallow merge where:
+    /// <list type="number">
+    /// <item><description>All entries from <paramref name="dict1"/> are copied to the new dictionary</description></item>
+    /// <item><description>Entries from <paramref name="dict2"/> are then added or overwrite existing keys</description></item>
+    /// <item><description>Null values in <paramref name="dict2"/> will overwrite existing keys with null</description></item>
+    /// </list>
+    /// </para>
+    /// <para>
+    /// Special cases:
+    /// <list type="bullet">
+    /// <item><description>If <paramref name="dict2"/> is null, returns a copy of <paramref name="dict1"/></description></item>
+    /// <item><description>Original dictionaries remain unmodified</description></item>
+    /// <item><description>Performs a shallow copy (values are not cloned)</description></item>
+    /// </list>
+    /// </para>
+    /// <example>
+    /// Basic merge with key override:
+    /// <code>
+    /// var dict1 = new Dictionary&lt;string, object&gt; { ["a"] = 1, ["b"] = 2 };
+    /// var dict2 = new Dictionary&lt;string, object&gt; { ["b"] = 99, ["c"] = 3 };
+    /// 
+    /// var merged = dict1.Merge(dict2);
+    /// // Result: { "a": 1, "b": 99, "c": 3 }
+    /// </code>
+    /// 
+    /// Handling null merge:
+    /// <code>
+    /// var merged = dict1.Merge(null);
+    /// // Returns copy of dict1
+    /// </code>
+    /// </example>
+    /// </remarks>
+    public static Dictionary<string, object> Merge(this IDictionary<string, object> dict1,
+        IDictionary<string, object>? dict2)
+    {
+        var result = new Dictionary<string, object>(dict1);
+        if(dict2  != null)
+        {
+            foreach (var val in dict2)
+            {
+                result[val.Key] = val.Value;
+            }
+        }
+
+        return result;
+    }
+}

src/Paramore.Brighter/Extensions/RequestContextExtensions.cs

@@ -0,0 +1,152 @@
+using System.Collections.Generic;
+
+namespace Paramore.Brighter.Extensions;
+
+/// <summary>
+/// Provides extension methods for <see cref="IRequestContext"/> to simplify access to common context bag values.
+/// </summary>
+/// <remarks>
+/// These extensions offer type-safe access to Brighter's reserved context bag values,
+/// handling type conversion and null checks internally. They are safe to call with null contexts.
+/// </remarks>
+public static class RequestContextExtensions
+{
+    /// <summary>
+    /// Retrieves the partition key from the request context bag.
+    /// </summary>
+    /// <param name="context">The request context (may be null)</param>
+    /// <returns>
+    /// The partition key if present and valid, otherwise <see cref="PartitionKey.Empty"/>.
+    /// </returns>
+    /// <remarks>
+    /// <para>
+    /// Handles two valid value types in the context bag:
+    /// <list type="bullet">
+    /// <item><description><see cref="string"/>: Converted to a <see cref="PartitionKey"/> instance</description></item>
+    /// <item><description><see cref="PartitionKey"/>: Returned directly</description></item>
+    /// </list>
+    /// </para>
+    /// <para>
+    /// Returns <see cref="PartitionKey.Empty"/> for:
+    /// <list type="bullet">
+    /// <item><description>Null context</description></item>
+    /// <item><description>Missing partition key entry</description></item>
+    /// <item><description>Unsupported value types</description></item>
+    /// </list>
+    /// </para>
+    /// <example>
+    /// Usage:
+    /// <code>
+    /// var partitionKey = requestContext.GetPartitionKey();
+    /// if (!partitionKey.IsEmpty)
+    /// {
+    ///     // Use partition key
+    /// }
+    /// </code>
+    /// </example>
+    /// </remarks>
+    public static PartitionKey GetPartitionKey(this IRequestContext? context)
+    {
+        if (context == null || !context.Bag.TryGetValue(RequestContextBagNames.PartitionKey, out var tmp))
+        {
+            return PartitionKey.Empty;
+        }
+
+        return tmp switch
+        {
+            string partitionKeyAsString => new PartitionKey(partitionKeyAsString),
+            PartitionKey partitionKey => partitionKey,
+            _ => PartitionKey.Empty
+        };
+    }
+
+    /// <summary>
+    /// Retrieves the dynamic headers dictionary from the request context bag.
+    /// </summary>
+    /// <param name="context">The request context (may be null)</param>
+    /// <returns>
+    /// The headers dictionary if present and valid, otherwise null.
+    /// </returns>
+    /// <remarks>
+    /// <para>
+    /// The value must be stored in the context bag as a <see cref="Dictionary{TKey, TValue}"/> 
+    /// where TKey is <see cref="string"/> and TValue is <see cref="object"/>.
+    /// </para>
+    /// <para>
+    /// Returns null for:
+    /// <list type="bullet">
+    /// <item><description>Null context</description></item>
+    /// <item><description>Missing headers entry</description></item>
+    /// <item><description>Type mismatch (not Dictionary&lt;string, object&gt;)</description></item>
+    /// </list>
+    /// </para>
+    /// <example>
+    /// Usage:
+    /// <code>
+    /// var headers = requestContext.GetHeaders();
+    /// if (headers != null)
+    /// {
+    ///     foreach (var header in headers)
+    ///     {
+    ///         // Process headers
+    ///     }
+    /// }
+    /// </code>
+    /// </example>
+    /// </remarks>
+    public static Dictionary<string, object>? GetHeaders(this IRequestContext? context)
+    {
+        if (context != null 
+            && context.Bag.TryGetValue(RequestContextBagNames.Headers, out var tmp)
+            && tmp is Dictionary<string, object> headers)
+        {
+            return headers;
+        }
+
+        return null;
+    }
+
+    /// <summary>
+    /// Retrieves CloudEvent additional properties from the request context bag.
+    /// </summary>
+    /// <param name="context">The request context (may be null)</param>
+    /// <returns>
+    /// The CloudEvent extensions dictionary if present and valid, otherwise null.
+    /// </returns>
+    /// <remarks>
+    /// <para>
+    /// The value must be stored in the context bag as a <see cref="Dictionary{TKey, TValue}"/> 
+    /// where TKey is <see cref="string"/> and TValue is <see cref="object"/>.
+    /// </para>
+    /// <para>
+    /// Returns null for:
+    /// <list type="bullet">
+    /// <item><description>Null context</description></item>
+    /// <item><description>Missing CloudEventsAdditionalProperties entry</description></item>
+    /// <item><description>Type mismatch (not Dictionary&lt;string, object&gt;)</description></item>
+    /// </list>
+    /// </para>
+    /// <example>
+    /// Usage:
+    /// <code>
+    /// var cloudEventProps = requestContext.GetCloudEventAdditionalProperties();
+    /// if (cloudEventProps != null)
+    /// {
+    ///     // Add to CloudEvent extensions
+    /// }
+    /// </code>
+    /// </example>
+    /// <seealso cref="RequestContextBagNames.CloudEventsAdditionalProperties"/>
+    /// </remarks>
+    public static Dictionary<string, object>? GetCloudEventAdditionalProperties(this IRequestContext? context)
+    {
+        if (context != null 
+            && context.Bag.TryGetValue(RequestContextBagNames.CloudEventsAdditionalProperties, out var tmp)
+            && tmp is Dictionary<string, object> cloudEventAdditionalProperties)
+        {
+            return cloudEventAdditionalProperties;
+        }
+
+        return null;
+    }
+}

src/Paramore.Brighter/IRequestContext.cs

@@ -23,7 +23,6 @@ THE SOFTWARE. */
 #endregion
 
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Diagnostics;
 using Paramore.Brighter.FeatureSwitch;
 using Polly.Registry;
@@ -63,7 +62,7 @@ public interface IRequestContext
         /// Gets the Feature Switches
         /// </summary>
         IAmAFeatureSwitchRegistry? FeatureSwitches { get; }
-
+        
         /// <summary>
         /// Create a new copy of the Request Context
         /// </summary>