Merge pull request confluentinc#36 from confluentinc/subscription_assignment

Add Consumer.Subscription() and Assignment() APIs

Author: edenhill

kafka/consumer.go

@@ -24,6 +24,11 @@ import (
 /*
 #include <stdlib.h>
 #include <librdkafka/rdkafka.h>
+
+
+static rd_kafka_topic_partition_t *_c_rdkafka_topic_partition_list_entry(rd_kafka_topic_partition_list_t *rktparlist, int idx) {
+   return idx < rktparlist->cnt ? &rktparlist->elems[idx] : NULL;
+}
 */
 import "C"
 
@@ -201,7 +206,7 @@ func (c *Consumer) CommitOffsets(offsets []TopicPartition) ([]TopicPartition, er
 // (through Assign() or implicitly through a self-rebalanced Subscribe()).
 // To set the starting offset it is preferred to use Assign() and provide
 // a starting offset for each partition.
-// 
+//
 // Returns an error on failure or nil otherwise.
 func (c *Consumer) Seek(partition TopicPartition, timeoutMs int) error {
 	rkt := c.handle.getRkt(*partition.Topic)
@@ -389,3 +394,39 @@ func (c *Consumer) GetMetadata(topic *string, allTopics bool, timeoutMs int) (*M
 func (c *Consumer) QueryWatermarkOffsets(topic string, partition int32, timeoutMs int) (low, high int64, err error) {
 	return queryWatermarkOffsets(c, topic, partition, timeoutMs)
 }
+
+// Subscription returns the current subscription as set by Subscribe()
+func (c *Consumer) Subscription() (topics []string, err error) {
+	var cTopics *C.rd_kafka_topic_partition_list_t
+
+	cErr := C.rd_kafka_subscription(c.handle.rk, &cTopics)
+	if cErr != C.RD_KAFKA_RESP_ERR_NO_ERROR {
+		return nil, newError(cErr)
+	}
+	defer C.rd_kafka_topic_partition_list_destroy(cTopics)
+
+	topicCnt := int(cTopics.cnt)
+	topics = make([]string, topicCnt)
+	for i := 0; i < topicCnt; i++ {
+		crktpar := C._c_rdkafka_topic_partition_list_entry(cTopics,
+			C.int(i))
+		topics[i] = C.GoString(crktpar.topic)
+	}
+
+	return topics, nil
+}
+
+// Assignment returns the current partition assignments
+func (c *Consumer) Assignment() (partitions []TopicPartition, err error) {
+	var cParts *C.rd_kafka_topic_partition_list_t
+
+	cErr := C.rd_kafka_assignment(c.handle.rk, &cParts)
+	if cErr != C.RD_KAFKA_RESP_ERR_NO_ERROR {
+		return nil, newError(cErr)
+	}
+	defer C.rd_kafka_topic_partition_list_destroy(cParts)
+
+	partitions = newTopicPartitionsFromCparts(cParts)
+
+	return partitions, nil
+}

kafka/consumer_test.go

@@ -17,6 +17,9 @@
 package kafka
 
 import (
+	"fmt"
+	"reflect"
+	"sort"
 	"testing"
 )
 
@@ -85,3 +88,75 @@ func TestConsumerAPIs(t *testing.T) {
 		t.Errorf("Close failed: %s", err)
 	}
 }
+
+func TestConsumerSubscription(t *testing.T) {
+	c, err := NewConsumer(&ConfigMap{"group.id": "gotest"})
+	if err != nil {
+		t.Fatalf("%s", err)
+	}
+
+	topics := []string{"gotest1", "gotest2", "gotest3"}
+	sort.Strings(topics)
+
+	err = c.SubscribeTopics(topics, nil)
+	if err != nil {
+		t.Fatalf("SubscribeTopics failed: %s", err)
+	}
+
+	subscription, err := c.Subscription()
+	if err != nil {
+		t.Fatalf("Subscription() failed: %s", err)
+	}
+
+	sort.Strings(subscription)
+
+	t.Logf("Compare Subscription %v to original list of topics %v\n",
+		subscription, topics)
+
+	r := reflect.DeepEqual(topics, subscription)
+	if r != true {
+		t.Fatalf("Subscription() %v does not match original topics %v",
+			subscription, topics)
+	}
+	c.Close()
+}
+
+func TestConsumerAssignment(t *testing.T) {
+	c, err := NewConsumer(&ConfigMap{"group.id": "gotest"})
+	if err != nil {
+		t.Fatalf("%s", err)
+	}
+
+	topic0 := "topic0"
+	topic1 := "topic1"
+	partitions := TopicPartitions{
+		{Topic: &topic1, Partition: 1},
+		{Topic: &topic1, Partition: 3},
+		{Topic: &topic0, Partition: 2}}
+	sort.Sort(partitions)
+
+	err = c.Assign(partitions)
+	if err != nil {
+		t.Fatalf("Assign failed: %s", err)
+	}
+
+	assignment, err := c.Assignment()
+	if err != nil {
+		t.Fatalf("Assignment() failed: %s", err)
+	}
+
+	sort.Sort(TopicPartitions(assignment))
+
+	t.Logf("Compare Assignment %v to original list of partitions %v\n",
+		assignment, partitions)
+
+	// reflect.DeepEqual() can't be used since TopicPartition.Topic
+	// is a pointer to a string rather than a string and the pointer
+	// will differ between partitions and assignment.
+	// Instead do a simple stringification + string compare.
+	if fmt.Sprintf("%v", assignment) != fmt.Sprintf("%v", partitions) {
+		t.Fatalf("Assignment() %v does not match original partitions %v",
+			assignment, partitions)
+	}
+	c.Close()
+}

kafka/kafka.go

@@ -264,6 +264,27 @@ func (p TopicPartition) String() string {
 		topic, p.Partition, p.Offset)
 }
 
+// TopicPartitions is a slice of TopicPartitions that also implements
+// the sort interface
+type TopicPartitions []TopicPartition
+
+func (tps TopicPartitions) Len() int {
+	return len(tps)
+}
+
+func (tps TopicPartitions) Less(i, j int) bool {
+	if *tps[i].Topic < *tps[j].Topic {
+		return true
+	} else if *tps[i].Topic > *tps[j].Topic {
+		return false
+	}
+	return tps[i].Partition < tps[j].Partition
+}
+
+func (tps TopicPartitions) Swap(i, j int) {
+	tps[i], tps[j] = tps[j], tps[i]
+}
+
 // new_cparts_from_TopicPartitions creates a new C rd_kafka_topic_partition_list_t
 // from a TopicPartition array.
 func newCPartsFromTopicPartitions(partitions []TopicPartition) (cparts *C.rd_kafka_topic_partition_list_t) {