From 8978f04e9de63ba18029db4999788f5e6225c743 Mon Sep 17 00:00:00 2001
From: shlomiassaf <shlomi.assaf@sharegain.com>
Date: Thu, 6 Jun 2024 22:59:05 +0300
Subject: [PATCH] feat: support dynamic configuration when creating topic

---
 .../IClusterConfigurationBuilder.cs           |  6 ++++-
 src/KafkaFlow/Clusters/ClusterManager.cs      |  2 ++
 .../ClusterConfigurationBuilder.cs            |  6 +++--
 .../Configuration/TopicConfiguration.cs       | 25 ++++++++++++++++++-
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/KafkaFlow.Abstractions/Configuration/IClusterConfigurationBuilder.cs b/src/KafkaFlow.Abstractions/Configuration/IClusterConfigurationBuilder.cs
index 2c6968baa..858fd058b 100644
--- a/src/KafkaFlow.Abstractions/Configuration/IClusterConfigurationBuilder.cs
+++ b/src/KafkaFlow.Abstractions/Configuration/IClusterConfigurationBuilder.cs
@@ -77,9 +77,13 @@ public interface IClusterConfigurationBuilder
     /// <param name="topicName">The topic name</param>
     /// <param name="numberOfPartitions">The number of Topic partitions. Default is to use the cluster-defined partitions.</param>
     /// <param name="replicationFactor">The Topic replication factor. Default is to use the cluster-defined replication factor.</param>
+    /// <param name="configs">The configuration to use to create the new topic</param>
+    /// <param name="replicasAssignments">A map from partition id to replica ids</param>
     /// <returns></returns>
     IClusterConfigurationBuilder CreateTopicIfNotExists(
         string topicName,
         int numberOfPartitions = -1,
-        short replicationFactor = -1);
+        short replicationFactor = -1,
+        IDictionary<string, string> configs = null,
+        IDictionary<int, IEnumerable<int>> replicasAssignments = null);
 }
diff --git a/src/KafkaFlow/Clusters/ClusterManager.cs b/src/KafkaFlow/Clusters/ClusterManager.cs
index 814088c50..dd37288cf 100644
--- a/src/KafkaFlow/Clusters/ClusterManager.cs
+++ b/src/KafkaFlow/Clusters/ClusterManager.cs
@@ -122,6 +122,8 @@ public async Task CreateIfNotExistsAsync(IEnumerable<TopicConfiguration> configu
                         Name = topicConfiguration.Name,
                         ReplicationFactor = topicConfiguration.Replicas,
                         NumPartitions = topicConfiguration.Partitions,
+                        Configs = topicConfiguration.Configs.ToDictionary(x => x.Key, x => x.Value),
+                        ReplicasAssignments = topicConfiguration.ReplicasAssignments.ToDictionary(x => x.Key, x => x.Value.ToList()),
                     })
                 .ToArray();
 
diff --git a/src/KafkaFlow/Configuration/ClusterConfigurationBuilder.cs b/src/KafkaFlow/Configuration/ClusterConfigurationBuilder.cs
index cd85129ef..e594c0520 100644
--- a/src/KafkaFlow/Configuration/ClusterConfigurationBuilder.cs
+++ b/src/KafkaFlow/Configuration/ClusterConfigurationBuilder.cs
@@ -111,9 +111,11 @@ public IClusterConfigurationBuilder OnStarted(Action<IDependencyResolver> handle
     public IClusterConfigurationBuilder CreateTopicIfNotExists(
         string topicName,
         int numberOfPartitions = -1,
-        short replicationFactor = -1)
+        short replicationFactor = -1,
+        IDictionary<string, string> configs = null,
+        IDictionary<int, IEnumerable<int>> replicasAssignments = null)
     {
-        _topicsToCreateIfNotExist.Add(new TopicConfiguration(topicName, numberOfPartitions, replicationFactor));
+        _topicsToCreateIfNotExist.Add(new TopicConfiguration(topicName, numberOfPartitions, replicationFactor, configs, replicasAssignments));
         return this;
     }
 }
diff --git a/src/KafkaFlow/Configuration/TopicConfiguration.cs b/src/KafkaFlow/Configuration/TopicConfiguration.cs
index f2eac31bb..672a559fe 100644
--- a/src/KafkaFlow/Configuration/TopicConfiguration.cs
+++ b/src/KafkaFlow/Configuration/TopicConfiguration.cs
@@ -1,3 +1,7 @@
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Linq;
+
 namespace KafkaFlow.Configuration;
 
 /// <summary>
@@ -5,17 +9,26 @@ namespace KafkaFlow.Configuration;
 /// </summary>
 public class TopicConfiguration
 {
+    private static readonly IReadOnlyDictionary<string, string> s_emptyConfigs = new ReadOnlyDictionary<string, string>(new Dictionary<string, string>());
+    private static readonly IReadOnlyDictionary<int, IReadOnlyList<int>> s_emptyReplicasAssignments = new ReadOnlyDictionary<int, IReadOnlyList<int>>(new Dictionary<int, IReadOnlyList<int>>());
+
     /// <summary>
     /// Initializes a new instance of the <see cref="TopicConfiguration"/> class.
     /// </summary>
     /// <param name="name">The topic name</param>
     /// <param name="partitions">The number of partitions for the topic</param>
     /// <param name="replicas">Replication factor for the topic</param>
-    public TopicConfiguration(string name, int partitions, short replicas)
+    /// <param name="configs">The configuration to use to create the new topic</param>
+    /// <param name="replicasAssignments">A map from partition id to replica ids</param>
+    public TopicConfiguration(string name, int partitions, short replicas, IDictionary<string, string> configs = null, IDictionary<int, IEnumerable<int>> replicasAssignments = null)
     {
         this.Name = name;
         this.Partitions = partitions;
         this.Replicas = replicas;
+        this.Configs = configs is null ? s_emptyConfigs : new ReadOnlyDictionary<string, string>(configs);
+        this.ReplicasAssignments = replicasAssignments is null
+            ? s_emptyReplicasAssignments
+            : new ReadOnlyDictionary<int, IReadOnlyList<int>>(replicasAssignments.ToDictionary(x => x.Key, x => (IReadOnlyList<int>)x.Value.ToList()));
     }
 
     /// <summary>
@@ -32,4 +45,14 @@ public TopicConfiguration(string name, int partitions, short replicas)
     /// Gets the Topic Replication Factor
     /// </summary>
     public short Replicas { get; }
+
+    /// <summary>Gets the configuration to use to create the new topic.</summary>
+    public IReadOnlyDictionary<string, string> Configs { get; }
+
+    /// <summary>
+    ///     Gets the mapping from partition id to replica ids (i.e., static broker ids) or null
+    ///     if the number of partitions and replication factor are specified
+    ///     instead.
+    /// </summary>
+    public IReadOnlyDictionary<int, IReadOnlyList<int>> ReplicasAssignments { get; }
 }