refactor: use same type as assignments for easier usage

Signed-off-by: Mirza Brunjadze <[email protected]>

Author: xxzefgh

src/clients/consumer/consumer.ts

@@ -99,7 +99,7 @@ import {
   type ListOffsetsOptions,
   type Offsets,
   type OffsetsWithTimestamps,
-  type TopicPartition
+  type TopicPartitions
 } from './types.ts'
 
 export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderValue = Buffer> extends Base<
@@ -300,67 +300,72 @@ export class Consumer<Key = Buffer, Value = Buffer, HeaderKey = Buffer, HeaderVa
     return callback![kCallbackPromise]
   }
 
-  pause (partitions: TopicPartition[]): void {
+  pause (topicPartitions: TopicPartitions[]): void {
     if (!this.assignments) {
       throw new UserError('Cannot pause partitions before joining a consumer group.')
     }
 
-    for (const { topic, partition } of partitions) {
+    for (const { topic, partitions } of topicPartitions) {
       const assignment = this.assignments.find(a => a.topic === topic)
       if (!assignment) {
         throw new UserError(`Topic '${topic}' is not assigned to this consumer.`, { topic })
       }
 
-      const existing = this.#pausedPartitions.get(topic)
-      if (existing) {
-        existing.add(partition)
-      } else {
-        this.#pausedPartitions.set(topic, new Set([partition]))
+      for (const partition of partitions) {
+        const existing = this.#pausedPartitions.get(topic)
+        if (existing) {
+          existing.add(partition)
+        } else {
+          this.#pausedPartitions.set(topic, new Set([partition]))
+        }
       }
     }
   }
 
-  resume (partitions: TopicPartition[]): void {
+  resume (topicPartitions: TopicPartitions[]): void {
     if (!this.assignments) {
       throw new UserError('Cannot resume partitions before joining a consumer group.')
     }
 
     let emitResumeEvent = false
-    for (const { topic, partition } of partitions) {
+    for (const { topic, partitions } of topicPartitions) {
       const assignment = this.assignments.find(a => a.topic === topic)
       if (!assignment) {
         throw new UserError(`Topic '${topic}' is not assigned to this consumer.`, { topic })
       }
 
-      const existing = this.#pausedPartitions.get(topic)
-      if (existing?.has(partition)) {
-        emitResumeEvent = true
-        existing.delete(partition)
-      }
+      for (const partition of partitions) {
+        const existing = this.#pausedPartitions.get(topic)
+        if (existing?.has(partition)) {
+          emitResumeEvent = true
+          existing.delete(partition)
+        }
 
-      if (existing?.size === 0) {
-        this.#pausedPartitions.delete(topic)
+        if (existing?.size === 0) {
+          this.#pausedPartitions.delete(topic)
+        }
       }
     }
 
     if (emitResumeEvent) {
-      this.emitWithDebug('consumer', 'user:resume', { partitions })
+      this.emitWithDebug('consumer', 'user:resume', { partitions: topicPartitions })
     }
   }
 
-  paused (): TopicPartition[] {
-    const result: TopicPartition[] = []
+  paused (): TopicPartitions[] {
+    const result: TopicPartitions[] = []
     for (const [topic, partitions] of this.#pausedPartitions.entries()) {
-      for (const partition of partitions) {
-        result.push({ topic, partition })
-      }
+      result.push({
+        topic,
+        partitions: Array.from(partitions)
+      })
     }
 
     return result
   }
 
-  isPaused (partition: TopicPartition): boolean {
-    return !!this.#pausedPartitions.get(partition.topic)?.has(partition.partition)
+  isPaused (topic: string, partition: number): boolean {
+    return !!this.#pausedPartitions.get(topic)?.has(partition)
   }
 
   fetch (options: FetchOptions<Key, Value, HeaderKey, HeaderValue>, callback: CallbackWithPromise<FetchResponse>): void

src/clients/consumer/messages-stream.ts

@@ -380,7 +380,7 @@ export class MessagesStream<Key, Value, HeaderKey, HeaderValue> extends Readable
         const partitions = assignment.partitions
 
         for (const partition of partitions) {
-          if (this.#consumer.isPaused({ topic, partition })) {
+          if (this.#consumer.isPaused(topic, partition)) {
             continue
           }
 

src/clients/consumer/types.ts

@@ -10,14 +10,13 @@ export interface GroupProtocolSubscription {
   metadata?: Buffer | string
 }
 
-export interface TopicPartition {
+export interface TopicPartitions {
   topic: string
-  partition: number
+  partitions: number[]
 }
 
-export interface GroupAssignment {
-  topic: string
-  partitions: number[]
+export interface GroupAssignment extends TopicPartitions {
+  // Pulic export, left to not break existing usage
 }
 
 export interface GroupPartitionsAssignments {

test/clients/consumer/consumer.test.ts

@@ -3273,23 +3273,22 @@ test('pause should prevent fetches from paused partitions during consumption', a
     value: `value-${i}`
   }))
 
-  await produceTestMessages({ t, messages: messages.slice(0, 5), overrideOptions: { partitioner: () => 0 } })
-  await produceTestMessages({ t, messages: messages.slice(5, 10), overrideOptions: { partitioner: () => 1 } })
-
   const consumer = createConsumer(t)
   const stream = await consumer.consume({
     topics: [topic],
     mode: MessagesStreamModes.EARLIEST,
     autocommit: true
   })
 
-  consumer.pause([{ topic, partition: 0 }])
+  consumer.pause([{ topic, partitions: [0] }])
+  await produceTestMessages({ t, messages: messages.slice(0, 5), overrideOptions: { partitioner: () => 0 } })
+  await produceTestMessages({ t, messages: messages.slice(5, 10), overrideOptions: { partitioner: () => 1 } })
 
-  const received: number[] = []
+  const received: { key: string; partition: number }[] = []
   let count = 0
 
   for await (const message of stream) {
-    received.push(message.partition)
+    received.push({ key: message.key.toString(), partition: message.partition })
     count++
 
     if (count === 5) {
@@ -3298,7 +3297,7 @@ test('pause should prevent fetches from paused partitions during consumption', a
   }
 
   strictEqual(
-    received.every(p => p === 1),
+    received.every(m => m.partition === 1),
     true
   )
 })
@@ -3322,7 +3321,7 @@ test('resume should allow fetches from previously paused partitions', async t =>
     maxWaitTime: 100
   })
 
-  consumer.pause([{ topic, partition: 0 }])
+  consumer.pause([{ topic, partitions: [0] }])
 
   const received: { key: string; partition: number }[] = []
   let count = 0
@@ -3332,7 +3331,7 @@ test('resume should allow fetches from previously paused partitions', async t =>
     count++
 
     if (count === 5) {
-      consumer.resume([{ topic, partition: 0 }])
+      consumer.resume([{ topic, partitions: [0] }])
     }
 
     if (count === 10) {
@@ -3357,23 +3356,23 @@ test('resume should handle resuming non-paused partitions gracefully', async t =
   consumer.topics.trackAll(topic)
   await consumer.joinGroup({})
 
-  doesNotThrow(() => consumer.resume([{ topic, partition: 0 }]))
+  doesNotThrow(() => consumer.resume([{ topic, partitions: [0] }]))
 })
 
 test('pause/resume should throw error if consumer has not joined a group', async t => {
   const topic = await createTopic(t, true)
   const consumer = createConsumer(t)
 
   try {
-    consumer.pause([{ topic, partition: 0 }])
+    consumer.pause([{ topic, partitions: [0] }])
     throw new Error('Expected error not thrown')
   } catch (error) {
     strictEqual(error instanceof UserError, true)
     strictEqual(error.message, 'Cannot pause partitions before joining a consumer group.')
   }
 
   try {
-    consumer.resume([{ topic, partition: 0 }])
+    consumer.resume([{ topic, partitions: [0] }])
     throw new Error('Expected error not thrown')
   } catch (error) {
     strictEqual(error instanceof UserError, true)
@@ -3390,15 +3389,15 @@ test('pause/resume should throw error if topic is not assigned to consumer', asy
   await consumer.joinGroup({})
 
   try {
-    consumer.pause([{ topic: topic2, partition: 0 }])
+    consumer.pause([{ topic: topic2, partitions: [0] }])
     throw new Error('Expected error not thrown')
   } catch (error) {
     strictEqual(error instanceof UserError, true)
     strictEqual(error.message, `Topic '${topic2}' is not assigned to this consumer.`)
   }
 
   try {
-    consumer.resume([{ topic: topic2, partition: 0 }])
+    consumer.resume([{ topic: topic2, partitions: [0] }])
     throw new Error('Expected error not thrown')
   } catch (error) {
     strictEqual(error instanceof UserError, true)
@@ -3415,20 +3414,20 @@ test('pause/resume should handle multiple topic-partitions', async t => {
   await consumer.joinGroup({})
 
   consumer.pause([
-    { topic: topic1, partition: 0 },
-    { topic: topic2, partition: 0 }
+    { topic: topic1, partitions: [0] },
+    { topic: topic2, partitions: [0] }
   ])
 
-  strictEqual(consumer.isPaused({ topic: topic1, partition: 0 }), true)
-  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), true)
+  strictEqual(consumer.isPaused(topic1, 0), true)
+  strictEqual(consumer.isPaused(topic2, 0), true)
 
   consumer.resume([
-    { topic: topic1, partition: 0 },
-    { topic: topic2, partition: 0 }
+    { topic: topic1, partitions: [0] },
+    { topic: topic2, partitions: [0] }
   ])
 
-  strictEqual(consumer.isPaused({ topic: topic1, partition: 0 }), false)
-  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), false)
+  strictEqual(consumer.isPaused(topic1, 0), false)
+  strictEqual(consumer.isPaused(topic2, 0), false)
 })
 
 test('paused should return all paused topic-partitions', async t => {
@@ -3440,15 +3439,13 @@ test('paused should return all paused topic-partitions', async t => {
   await consumer.joinGroup({})
 
   consumer.pause([
-    { topic: topic1, partition: 0 },
-    { topic: topic1, partition: 1 },
-    { topic: topic2, partition: 0 }
+    { topic: topic1, partitions: [0, 1] },
+    { topic: topic2, partitions: [0] }
   ])
 
   deepStrictEqual(consumer.paused(), [
-    { topic: topic1, partition: 0 },
-    { topic: topic1, partition: 1 },
-    { topic: topic2, partition: 0 }
+    { topic: topic1, partitions: [0, 1] },
+    { topic: topic2, partitions: [0] }
   ])
 })
 
@@ -3459,8 +3456,8 @@ test('isPaused should return true for paused topic-partitions', async t => {
   consumer.topics.trackAll(topic)
   await consumer.joinGroup({})
 
-  consumer.pause([{ topic, partition: 0 }])
-  strictEqual(consumer.isPaused({ topic, partition: 0 }), true)
+  consumer.pause([{ topic, partitions: [0] }])
+  strictEqual(consumer.isPaused(topic, 0), true)
 })
 
 test('isPaused should return false for non-paused topic-partitions', async t => {
@@ -3470,7 +3467,7 @@ test('isPaused should return false for non-paused topic-partitions', async t =>
   consumer.topics.trackAll(topic)
   await consumer.joinGroup({})
 
-  strictEqual(consumer.isPaused({ topic, partition: 0 }), false)
+  strictEqual(consumer.isPaused(topic, 0), false)
 })
 
 test('isPaused should return false for topic-partitions not assigned to consumer', async t => {
@@ -3481,5 +3478,5 @@ test('isPaused should return false for topic-partitions not assigned to consumer
   consumer.topics.trackAll(topic1)
   await consumer.joinGroup({})
 
-  strictEqual(consumer.isPaused({ topic: topic2, partition: 0 }), false)
+  strictEqual(consumer.isPaused(topic2, 0), false)
 })