feat: Make Kafka group.id configurable. Fixes #89 (#464)

Signed-off-by: Alex Collins <[email protected]>
argoproj-labs · Oct 25, 2021 · 5ec67f6 · 5ec67f6
1 parent 782c39a
commit 5ec67f6
Show file tree

Hide file tree

Showing 15 changed files with 264 additions and 151 deletions.
diff --git a/api/v1alpha1/generated.pb.go b/api/v1alpha1/generated.pb.go
diff --git a/api/v1alpha1/generated.proto b/api/v1alpha1/generated.proto
diff --git a/api/v1alpha1/kafka_source.go b/api/v1alpha1/kafka_source.go
@@ -13,6 +13,8 @@ type KafkaSource struct {
 	FetchMin *resource.Quantity `json:"fetchMin,omitempty" protobuf:"bytes,3,opt,name=fetchMin"`
 	// +kubebuilder:default="500ms"
 	FetchWaitMax *metav1.Duration `json:"fetchWaitMax,omitempty" protobuf:"bytes,4,opt,name=fetchWaitMax"`
+	// GroupID is the consumer group ID. If not specified, a unique deterministic group ID is generated.
+	GroupID string `json:"groupId,omitempty" protobuf:"bytes,5,opt,name=groupId"`
 }
 
 func (m *KafkaSource) GetAutoOffsetReset() string {
@@ -26,3 +28,10 @@ func (m *KafkaSource) GetFetchMinBytes() int {
 func (m *KafkaSource) GetFetchWaitMaxMs() int {
 	return int(m.FetchWaitMax.Milliseconds())
 }
+
+func (m *KafkaSource) GetGroupID(defaultGroupID string) string {
+	if m.GroupID != "" {
+		return m.GroupID
+	}
+	return defaultGroupID
+}
diff --git a/api/v1alpha1/kafka_source_test.go b/api/v1alpha1/kafka_source_test.go
@@ -24,3 +24,14 @@ func TestKafkaSource_GetFetchWaitMaxMs(t *testing.T) {
 	s := KafkaSource{FetchWaitMax: &metav1.Duration{Duration: time.Second}}
 	assert.Equal(t, 1000, s.GetFetchWaitMaxMs())
 }
+
+func TestKafkaSource_GetGroupID(t *testing.T) {
+	t.Run("Default", func(t *testing.T) {
+		s := &KafkaSource{}
+		assert.Equal(t, "foo", s.GetGroupID("foo"))
+	})
+	t.Run("Specified", func(t *testing.T) {
+		s := &KafkaSource{GroupID: "bar"}
+		assert.Equal(t, "bar", s.GetGroupID("foo"))
+	})
+}
diff --git a/config/ci.yaml b/config/ci.yaml
@@ -3646,6 +3646,11 @@ spec:
                               fetchWaitMax:
                                 default: 500ms
                                 type: string
+                              groupId:
+                                description: GroupID is the consumer group ID. If
+                                  not specified, a unique deterministic group ID is
+                                  generated.
+                                type: string
                               maxMessageBytes:
                                 format: int32
                                 type: integer
@@ -10689,6 +10694,10 @@ spec:
                         fetchWaitMax:
                           default: 500ms
                           type: string
+                        groupId:
+                          description: GroupID is the consumer group ID. If not specified,
+                            a unique deterministic group ID is generated.
+                          type: string
                         maxMessageBytes:
                           format: int32
                           type: integer

diff --git a/config/crd/bases/dataflow.argoproj.io_pipelines.yaml b/config/crd/bases/dataflow.argoproj.io_pipelines.yaml
@@ -3641,6 +3641,11 @@ spec:
                               fetchWaitMax:
                                 default: 500ms
                                 type: string
+                              groupId:
+                                description: GroupID is the consumer group ID. If
+                                  not specified, a unique deterministic group ID is
+                                  generated.
+                                type: string
                               maxMessageBytes:
                                 format: int32
                                 type: integer

diff --git a/config/crd/bases/dataflow.argoproj.io_steps.yaml b/config/crd/bases/dataflow.argoproj.io_steps.yaml
@@ -3468,6 +3468,10 @@ spec:
                         fetchWaitMax:
                           default: 500ms
                           type: string
+                        groupId:
+                          description: GroupID is the consumer group ID. If not specified,
+                            a unique deterministic group ID is generated.
+                          type: string
                         maxMessageBytes:
                           format: int32
                           type: integer

diff --git a/config/default.yaml b/config/default.yaml
@@ -3646,6 +3646,11 @@ spec:
                               fetchWaitMax:
                                 default: 500ms
                                 type: string
+                              groupId:
+                                description: GroupID is the consumer group ID. If
+                                  not specified, a unique deterministic group ID is
+                                  generated.
+                                type: string
                               maxMessageBytes:
                                 format: int32
                                 type: integer
@@ -10689,6 +10694,10 @@ spec:
                         fetchWaitMax:
                           default: 500ms
                           type: string
+                        groupId:
+                          description: GroupID is the consumer group ID. If not specified,
+                            a unique deterministic group ID is generated.
+                          type: string
                         maxMessageBytes:
                           format: int32
                           type: integer

diff --git a/config/dev.yaml b/config/dev.yaml
@@ -3646,6 +3646,11 @@ spec:
                               fetchWaitMax:
                                 default: 500ms
                                 type: string
+                              groupId:
+                                description: GroupID is the consumer group ID. If
+                                  not specified, a unique deterministic group ID is
+                                  generated.
+                                type: string
                               maxMessageBytes:
                                 format: int32
                                 type: integer
@@ -10689,6 +10694,10 @@ spec:
                         fetchWaitMax:
                           default: 500ms
                           type: string
+                        groupId:
+                          description: GroupID is the consumer group ID. If not specified,
+                            a unique deterministic group ID is generated.
+                          type: string
                         maxMessageBytes:
                           format: int32
                           type: integer

diff --git a/config/quick-start.yaml b/config/quick-start.yaml
@@ -3646,6 +3646,11 @@ spec:
                               fetchWaitMax:
                                 default: 500ms
                                 type: string
+                              groupId:
+                                description: GroupID is the consumer group ID. If
+                                  not specified, a unique deterministic group ID is
+                                  generated.
+                                type: string
                               maxMessageBytes:
                                 format: int32
                                 type: integer
@@ -10689,6 +10694,10 @@ spec:
                         fetchWaitMax:
                           default: 500ms
                           type: string
+                        groupId:
+                          description: GroupID is the consumer group ID. If not specified,
+                            a unique deterministic group ID is generated.
+                          type: string
                         maxMessageBytes:
                           format: int32
                           type: integer

diff --git a/dsls/python/argo_dataflow/pipeline.py b/dsls/python/argo_dataflow/pipeline.py
@@ -600,13 +600,14 @@ def dump(self):
 
 
 class KafkaSource(Source):
-    def __init__(self, topic, name=None, retry=None, startOffset=None, fetchMin=None, fetchWaitMax=None):
+    def __init__(self, topic, name=None, retry=None, startOffset=None, fetchMin=None, fetchWaitMax=None, groupId=None):
         super().__init__(name=name, retry=retry)
         assert topic
         self._topic = topic
         self._startOffset = startOffset
         self._fetchMin = fetchMin
         self._fetchWaitMax = fetchWaitMax
+        self._groupId = groupId
 
     def dump(self):
         x = super().dump()
@@ -617,6 +618,8 @@ def dump(self):
             y["fetchMin"] = self._fetchMin
         if self._fetchWaitMax:
             y["fetchWaitMax"] = self._fetchWaitMax
+        if self._groupId:
+            y["groupId"] = self._groupId
         x['kafka'] = y
         return x
 
@@ -655,9 +658,9 @@ def http(name=None, retry=None, serviceName=None):
     return HTTPSource(name=name, serviceName=serviceName, retry=retry)
 
 
-def kafka(topic=None, name=None, retry=None, startOffset=None, fetchMin=None, fetchWaitMax=None):
+def kafka(topic=None, name=None, retry=None, startOffset=None, fetchMin=None, fetchWaitMax=None, groupId=None):
     return KafkaSource(topic, name=name, retry=retry, startOffset=startOffset, fetchMin=fetchMin,
-                       fetchWaitMax=fetchWaitMax)
+                       fetchWaitMax=fetchWaitMax, groupId=groupId)
 
 
 def stan(subject=None, name=None, retry=None):

diff --git a/examples/301-kafka-pipeline.py b/examples/301-kafka-pipeline.py
@@ -11,7 +11,7 @@
      """)
      .annotate("dataflow.argoproj.io/test", "true")
      .step(
-        (kafka('input-topic')
+        (kafka('input-topic', groupId='my-group')
          .cat()
          .kafka('output-topic', a_sync=True)
          ))

diff --git a/examples/301-kafka-pipeline.yaml b/examples/301-kafka-pipeline.yaml
@@ -20,4 +20,5 @@ spec:
         topic: output-topic
     sources:
     - kafka:
+        groupId: my-group
         topic: input-topic
diff --git a/runner/sidecar/source/kafka/kafka.go b/runner/sidecar/source/kafka/kafka.go
@@ -37,13 +37,14 @@ const (
 	pendingUnavailable = math.MinInt32
 )
 
-func New(ctx context.Context, secretInterface corev1.SecretInterface, consumerGroupID, sourceName, sourceURN string, replica int, x dfv1.KafkaSource, process source.Process) (source.Interface, error) {
+func New(ctx context.Context, secretInterface corev1.SecretInterface, cluster, namespace, pipelineName, stepName, sourceName, sourceURN string, replica int, x dfv1.KafkaSource, process source.Process) (source.Interface, error) {
 	logger := sharedutil.NewLogger().WithValues("source", sourceName)
 	config, err := sharedkafka.GetConfig(ctx, secretInterface, x.KafkaConfig)
 	if err != nil {
 		return nil, err
 	}
 	config["go.logs.channel.enable"] = true
+	consumerGroupID := x.GetGroupID(sharedutil.GetSourceUID(cluster, namespace, pipelineName, stepName, sourceName))
 	config["group.id"] = consumerGroupID
 	config["group.instance.id"] = fmt.Sprintf("%s/%d", consumerGroupID, replica)
 	config["heartbeat.interval.ms"] = 3 * seconds

diff --git a/runner/sidecar/sources.go b/runner/sidecar/sources.go
@@ -150,8 +150,7 @@ func connectSources(ctx context.Context, process func(context.Context, []byte) e
 				sources[sourceName] = y
 			}
 		} else if x := s.Kafka; x != nil {
-			groupID := sharedutil.GetSourceUID(cluster, namespace, pipelineName, stepName, sourceName)
-			if y, err := kafkasource.New(ctx, secretInterface, groupID, sourceName, sourceURN, replica, *x, processWithRetry); err != nil {
+			if y, err := kafkasource.New(ctx, secretInterface, cluster, namespace, pipelineName, stepName, sourceName, sourceURN, replica, *x, processWithRetry); err != nil {
 				return err
 			} else {
 				sources[sourceName] = y