Merge pull request #38 from vahidhashemian/add_internal_subscribe_to_…

…memq_consumer_assign Add an internal subscribe call to MemQ consumer's assign()
pinterest · Jun 20, 2024 · 6da931c · 6da931c
2 parents 7c90b23 + e1d1936
commit 6da931c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/psc-integration-test/src/test/java/com/pinterest/psc/consumer/memq/TestPscMemqConsumer.java b/psc-integration-test/src/test/java/com/pinterest/psc/consumer/memq/TestPscMemqConsumer.java
@@ -168,6 +168,21 @@ void unsubscribe() throws Exception {
         pscMemqConsumer.close();
     }
 
+    @Test
+    void testAssignDoesNotRequireSubscribe() throws Exception {
+        PscMemqConsumer<byte[], byte[]> pscMemqConsumer = getPscMemqConsumer(
+                "testAssignDoesNotRequireSubscribe");
+
+        pscMemqConsumer.assign(Collections.emptySet());
+
+        MemqTopicUri uri1 = MemqTopicUri.validate(TopicUri.validate(testMemqTopic1));
+        TopicUriPartition topicUriPartition = TestUtils.getFinalizedTopicUriPartition(uri1, 0);
+
+        pscMemqConsumer.assign(Sets.newHashSet(topicUriPartition));
+
+        pscMemqConsumer.close();
+    }
+
     @Test
     void close() throws Exception {
         PscMemqConsumer<byte[], byte[]> pscMemqConsumer = getPscMemqConsumer(

diff --git a/psc/src/main/java/com/pinterest/psc/consumer/memq/PscMemqConsumer.java b/psc/src/main/java/com/pinterest/psc/consumer/memq/PscMemqConsumer.java
@@ -181,6 +181,8 @@ public void assign(Set<TopicUriPartition> topicUriPartitions) throws ConsumerExc
         currentAssignment.retainAll(topicUriPartitions);
         currentAssignment.addAll(topicUriPartitions);
         try {
+            memqConsumer.subscribe(topicUriPartitions.stream().map(TopicUriPartition::getTopicUri)
+                    .map(TopicUri::getTopic).collect(Collectors.toSet()));
             memqConsumer.assign(topicUriPartitions.stream().map(TopicUriPartition::getPartition)
                     .collect(Collectors.toList()));
         } catch (Exception exception) {