Repository: nestjs/nest
Pull request #9954 (2022-07-19 08:56:49)

feat(microservices): add Kafka heartbeat callback to KafkaContext
## PR Checklist
Please check if your PR fulfills the following requirements:

- [x] The commit message follows our guidelines: https://github.com/nestjs/nest/blob/master/CONTRIBUTING.md
- [ ] Tests for the changes have been added (for bug fixes / features)
- [ ] Docs have been added / updated (for bug fixes / features)


## PR Type
Implements https://github.com/nestjs/nest/issues/9949

<!-- Please check the one that applies to this PR using "x". -->
- [ ] Bugfix
- [x] Feature
- [ ] Code style update (formatting, local variables)
- [ ] Refactoring (no functional changes, no api changes)
- [ ] Build related changes
- [ ] CI related changes
- [ ] Other... Please describe:

## What is the current behavior?
https://github.com/nestjs/nest/issues/9949

Issue Number: 9949


## What is the new behavior?
Adds heartbeat callback to KafkaContext (so it will be passed to the endpoint)

## Does this PR introduce a breaking change?
- [ ] Yes
- [x] No

<!-- If this PR contains a breaking change, please describe the impact and migration path for existing applications below. -->


## Other information


Patch:
diff --git a/packages/microservices/ctx-host/kafka.context.ts b/packages/microservices/ctx-host/kafka.context.ts
index 54c80475a..bd0969d94 100644
--- a/packages/microservices/ctx-host/kafka.context.ts
+++ b/packages/microservices/ctx-host/kafka.context.ts
@@ -6,6 +6,7 @@ type KafkaContextArgs = [
   partition: number,
   topic: string,
   consumer: Consumer,
+  heartbeat: () => Promise<void>,
 ];
 
 export class KafkaContext extends BaseRpcContext<KafkaContextArgs> {
@@ -40,4 +41,11 @@ export class KafkaContext extends BaseRpcContext<KafkaContextArgs> {
   getConsumer() {
     return this.args[3];
   }
+
+  /**
+   * Returns the Kafka heartbeat callback.
+   */
+   getHeartbeat() {
+    return this.args[4];
+  } 
 }
diff --git a/packages/microservices/external/kafka.interface.ts b/packages/microservices/external/kafka.interface.ts
index 164d11155..308981b48 100644
--- a/packages/microservices/external/kafka.interface.ts
+++ b/packages/microservices/external/kafka.interface.ts
@@ -870,6 +870,8 @@ export interface EachMessagePayload {
   topic: string;
   partition: number;
   message: KafkaMessage;
+  heartbeat(): Promise<void>;
+  pause(): () => void;
 }
 
 export interface EachBatchPayload {
diff --git a/packages/microservices/server/server-kafka.ts b/packages/microservices/server/server-kafka.ts
index 2a426eaf6..e7bfee7ce 100644
--- a/packages/microservices/server/server-kafka.ts
+++ b/packages/microservices/server/server-kafka.ts
@@ -167,6 +167,7 @@ export class ServerKafka extends Server implements CustomTransportStrategy {
       payload.partition,
       payload.topic,
       this.consumer,
+      payload.heartbeat,
     ]);
     const handler = this.getHandlerByPattern(packet.pattern);
     // if the correlation id or reply topic is not set


Test patch:
diff --git a/packages/microservices/test/client/client-kafka.spec.ts b/packages/microservices/test/client/client-kafka.spec.ts
index d68b8c28a..3770baea9 100644
--- a/packages/microservices/test/client/client-kafka.spec.ts
+++ b/packages/microservices/test/client/client-kafka.spec.ts
@@ -22,6 +22,8 @@ describe('ClientKafka', () => {
   const timestamp = new Date().toISOString();
   const attributes = 1;
   const messageValue = 'test-message';
+  const heartbeat = async () => {};
+  const pause = () => () => {};
 
   // message
   const message: KafkaMessage = {
@@ -57,6 +59,8 @@ describe('ClientKafka', () => {
       },
       message,
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadDisposed: EachMessagePayload = {
@@ -75,6 +79,8 @@ describe('ClientKafka', () => {
         value: { test: true },
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadError: EachMessagePayload = {
@@ -93,6 +99,8 @@ describe('ClientKafka', () => {
         value: null,
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const payloadWithoutCorrelation: EachMessagePayload = {
@@ -104,6 +112,8 @@ describe('ClientKafka', () => {
       },
       message,
     ),
+    heartbeat,
+    pause,
   };
 
   // deserialized payload
@@ -118,6 +128,8 @@ describe('ClientKafka', () => {
       },
       deserializedMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const deserializedPayloadDisposed: EachMessagePayload = {
@@ -136,6 +148,8 @@ describe('ClientKafka', () => {
         value: { test: true },
       },
     ),
+    heartbeat,
+    pause,
   };
 
   const deserializedPayloadError: EachMessagePayload = {
@@ -154,6 +168,8 @@ describe('ClientKafka', () => {
         value: null,
       },
     ),
+    heartbeat,
+    pause,
   };
 
   let client: ClientKafka;
diff --git a/packages/microservices/test/ctx-host/kafka.context.spec.ts b/packages/microservices/test/ctx-host/kafka.context.spec.ts
index 5f350a635..26db4a868 100644
--- a/packages/microservices/test/ctx-host/kafka.context.spec.ts
+++ b/packages/microservices/test/ctx-host/kafka.context.spec.ts
@@ -3,12 +3,12 @@ import { KafkaContext } from '../../ctx-host';
 import { Consumer, KafkaMessage } from '../../external/kafka.interface';
 
 describe('KafkaContext', () => {
-  const args = ['test', { test: true }, undefined, { test: 'consumer' }];
+  const args = ['test', { test: true }, undefined, { test: 'consumer' }, () => {}];
   let context: KafkaContext;
 
   beforeEach(() => {
     context = new KafkaContext(
-      args as [KafkaMessage, number, string, Consumer],
+      args as [KafkaMessage, number, string, Consumer, () => Promise<void>],
     );
   });
   describe('getTopic', () => {
@@ -31,4 +31,9 @@ describe('KafkaContext', () => {
       expect(context.getConsumer()).to.deep.eq({ test: 'consumer' });
     });
   });
+  describe('getHeartbeat', () => {
+    it('should return heartbeat callback', () => {
+      expect(context.getHeartbeat()).to.be.eql(args[4]);
+    });
+  });
 });
diff --git a/packages/microservices/test/server/server-kafka.spec.ts b/packages/microservices/test/server/server-kafka.spec.ts
index da129d56d..5c32118e4 100644
--- a/packages/microservices/test/server/server-kafka.spec.ts
+++ b/packages/microservices/test/server/server-kafka.spec.ts
@@ -26,6 +26,8 @@ describe('ServerKafka', () => {
   const key = '1';
   const timestamp = new Date().toISOString();
   const messageValue = 'test-message';
+  const heartbeat = async () => {};
+  const pause = () => () => {};
 
   const eventMessage: KafkaMessage = {
     key: Buffer.from(key),
@@ -44,6 +46,8 @@ describe('ServerKafka', () => {
       },
       eventMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const eventWithCorrelationIdPayload: EachMessagePayload = {
@@ -57,6 +61,8 @@ describe('ServerKafka', () => {
       },
       eventMessage,
     ),
+    heartbeat,
+    pause,
   };
 
   const message: KafkaMessage = Object.assign(
@@ -73,6 +79,8 @@ describe('ServerKafka', () => {
     topic,
     partition: 0,
     message,
+    heartbeat,
+    pause,
   };
 
   let server: ServerKafka;