Feat: Kafka execution context support (#292)

CarlosGamero · web-flow · commit 2eb8bbc948af · 2025-06-13T20:16:21.000+02:00
* Handler container supporting execution context

* Test fixes

* Adding execution context to the consumer

* Lint fix

* Release prepare
diff --git a/packages/kafka/lib/AbstractKafkaConsumer.ts b/packages/kafka/lib/AbstractKafkaConsumer.ts
@@ -24,13 +24,16 @@ import { safeJsonDeserializer } from './utils/safeJsonDeserializer.js'
 export type KafkaConsumerDependencies = KafkaDependencies &
   Pick<QueueConsumerDependencies, 'transactionObservabilityManager'>
 
-export type KafkaConsumerOptions<TopicsConfig extends TopicConfig[]> = BaseKafkaOptions &
+export type KafkaConsumerOptions<
+  TopicsConfig extends TopicConfig[],
+  ExecutionContext,
+> = BaseKafkaOptions &
   Omit<
     ConsumerOptions<string, object, string, string>,
     'deserializers' | 'autocommit' | keyof KafkaConfig
   > &
   Omit<ConsumeOptions<string, object, string, string>, 'topics'> & {
-    handlers: KafkaHandlerRouting<TopicsConfig>
+    handlers: KafkaHandlerRouting<TopicsConfig, ExecutionContext>
   }
 
 /*
@@ -41,24 +44,28 @@ const MAX_IN_MEMORY_RETRIES = 3
 
 export abstract class AbstractKafkaConsumer<
   TopicsConfig extends TopicConfig[],
-> extends AbstractKafkaService<TopicsConfig, KafkaConsumerOptions<TopicsConfig>> {
+  ExecutionContext,
+> extends AbstractKafkaService<TopicsConfig, KafkaConsumerOptions<TopicsConfig, ExecutionContext>> {
   private readonly consumer: Consumer<string, object, string, string>
   private consumerStream?: MessagesStream<string, object, string, string>
 
   private readonly transactionObservabilityManager: TransactionObservabilityManager
-  private readonly handlerContainer: KafkaHandlerContainer<TopicsConfig>
+  private readonly handlerContainer: KafkaHandlerContainer<TopicsConfig, ExecutionContext>
+  private readonly executionContext: ExecutionContext
 
   constructor(
     dependencies: KafkaConsumerDependencies,
-    options: KafkaConsumerOptions<TopicsConfig>,
+    options: KafkaConsumerOptions<TopicsConfig, ExecutionContext>,
+    executionContext: ExecutionContext,
   ) {
     super(dependencies, options)
 
     this.transactionObservabilityManager = dependencies.transactionObservabilityManager
-    this.handlerContainer = new KafkaHandlerContainer<TopicsConfig>(
+    this.handlerContainer = new KafkaHandlerContainer<TopicsConfig, ExecutionContext>(
       options.handlers,
       options.messageTypeField,
     )
+    this.executionContext = executionContext
 
     this.consumer = new Consumer({
       ...this.options.kafka,
@@ -169,11 +176,11 @@ export abstract class AbstractKafkaConsumer<
 
   private async tryToConsume<MessageValue extends object>(
     message: Message<string, MessageValue, string, string>,
-    handler: KafkaHandler<MessageValue>,
+    handler: KafkaHandler<MessageValue, ExecutionContext>,
     requestContext: RequestContext,
   ): Promise<boolean> {
     try {
-      await handler(message, requestContext)
+      await handler(message, this.executionContext, requestContext)
       return true
     } catch (error) {
       this.handlerError(error, {
diff --git a/packages/kafka/lib/handler-container/KafkaHandlerConfig.ts b/packages/kafka/lib/handler-container/KafkaHandlerConfig.ts
@@ -7,16 +7,20 @@ export interface RequestContext {
   reqId: string
 }
 
-export type KafkaHandler<MessageValue extends object> = (
+export type KafkaHandler<MessageValue extends object, ExecutionContext> = (
   message: Message<string, MessageValue, string, string>,
+  context: ExecutionContext,
   requestContext: RequestContext,
 ) => Promise<void> | void
 
-export class KafkaHandlerConfig<MessageValue extends object> {
+export class KafkaHandlerConfig<MessageValue extends object, ExecutionContext> {
   public readonly schema: ZodSchema<MessageValue>
-  public readonly handler: KafkaHandler<MessageValue>
+  public readonly handler: KafkaHandler<MessageValue, ExecutionContext>
 
-  constructor(schema: ZodSchema<MessageValue>, handler: KafkaHandler<MessageValue>) {
+  constructor(
+    schema: ZodSchema<MessageValue>,
+    handler: KafkaHandler<MessageValue, ExecutionContext>,
+  ) {
     this.schema = schema
     this.handler = handler
   }
diff --git a/packages/kafka/lib/handler-container/KafkaHandlerContainer.spec.ts b/packages/kafka/lib/handler-container/KafkaHandlerContainer.spec.ts
@@ -37,7 +37,7 @@ describe('KafkaHandlerContainer', () => {
           new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()),
           new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()),
         ],
-      } satisfies KafkaHandlerRouting<typeof topicsConfig>
+      }
 
       // When & Then
       expect(
@@ -50,7 +50,7 @@ describe('KafkaHandlerContainer', () => {
 
     it('should resolve handler with message type', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any> = {
+      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
         all: [
           new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()),
           new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()),
@@ -61,7 +61,7 @@ describe('KafkaHandlerContainer', () => {
       }
 
       // When
-      const container = new KafkaHandlerContainer<TopicsConfig>(topicHandlers, 'type')
+      const container = new KafkaHandlerContainer(topicHandlers, 'type')
 
       // Then
       expect(container.resolveHandler('all', { type: 'create' })?.schema).toBe(CREATE_SCHEMA)
@@ -79,7 +79,7 @@ describe('KafkaHandlerContainer', () => {
 
     it('should resolve handler without message type', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any> = {
+      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
         create: [new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve())],
         empty: [new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve())],
       }
@@ -102,15 +102,15 @@ describe('KafkaHandlerContainer', () => {
   describe('topics', () => {
     it('should not fail with empty topics', () => {
       // When
-      const container = new KafkaHandlerContainer<TopicsConfig>({})
+      const container = new KafkaHandlerContainer({})
 
       // Then
       expect(container.topics).toEqual([])
     })
 
     it('should return all topics', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any> = {
+      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
         all: [
           new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()),
           new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()),
@@ -122,7 +122,7 @@ describe('KafkaHandlerContainer', () => {
       }
 
       // When
-      const container = new KafkaHandlerContainer<TopicsConfig>(topicHandlers, 'type')
+      const container = new KafkaHandlerContainer(topicHandlers, 'type')
 
       // Then
       expect(container.topics).toEqual(['all', 'create', 'empty'])
diff --git a/packages/kafka/lib/handler-container/KafkaHandlerContainer.ts b/packages/kafka/lib/handler-container/KafkaHandlerContainer.ts
@@ -9,24 +9,30 @@ import type { KafkaHandlerRouting } from './KafkaHandlerRoutingBuilder.ts'
 
 const DEFAULT_HANDLER_KEY = Symbol('default-handler')
 
-type Handlers<TopicsConfig extends TopicConfig[]> = Record<
+type Handlers<TopicsConfig extends TopicConfig[], ExecutionContext> = Record<
   string,
-  Record<string | symbol, KafkaHandlerConfig<SupportedMessageValues<TopicsConfig>>>
+  Record<
+    string | symbol,
+    KafkaHandlerConfig<SupportedMessageValues<TopicsConfig>, ExecutionContext>
+  >
 >
 
-export class KafkaHandlerContainer<TopicsConfig extends TopicConfig[]> {
-  private readonly handlers: Handlers<TopicsConfig>
+export class KafkaHandlerContainer<TopicsConfig extends TopicConfig[], ExecutionContext> {
+  private readonly handlers: Handlers<TopicsConfig, ExecutionContext>
   private readonly messageTypeField?: string
 
-  constructor(topicHandlers: KafkaHandlerRouting<TopicsConfig>, messageTypeField?: string) {
+  constructor(
+    topicHandlers: KafkaHandlerRouting<TopicsConfig, ExecutionContext>,
+    messageTypeField?: string,
+  ) {
     this.messageTypeField = messageTypeField
     this.handlers = this.mapTopicHandlers(topicHandlers)
   }
 
   private mapTopicHandlers(
-    topicHandlerRouting: KafkaHandlerRouting<TopicsConfig>,
-  ): Handlers<TopicsConfig> {
-    const result: Handlers<TopicsConfig> = {}
+    topicHandlerRouting: KafkaHandlerRouting<TopicsConfig, ExecutionContext>,
+  ): Handlers<TopicsConfig, ExecutionContext> {
+    const result: Handlers<TopicsConfig, ExecutionContext> = {}
 
     for (const [topic, topicHandlers] of Object.entries(topicHandlerRouting)) {
       if (!topicHandlers.length) continue
@@ -52,7 +58,9 @@ export class KafkaHandlerContainer<TopicsConfig extends TopicConfig[]> {
   resolveHandler<Topic extends SupportedTopics<TopicsConfig>>(
     topic: Topic,
     messageValue: SupportedMessageValuesForTopic<TopicsConfig, Topic>,
-  ): KafkaHandlerConfig<SupportedMessageValuesForTopic<TopicsConfig, Topic>> | undefined {
+  ):
+    | KafkaHandlerConfig<SupportedMessageValuesForTopic<TopicsConfig, Topic>, ExecutionContext>
+    | undefined {
     const handlers = this.handlers[topic]
     if (!handlers) return undefined
 
diff --git a/packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.spec.ts b/packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.spec.ts
@@ -15,30 +15,37 @@ const topicsConfig = [
 ] as const satisfies TopicConfig[]
 type TopicsConfig = typeof topicsConfig
 
+type ExecutionContext = {
+  hello: string
+}
+
 describe('KafkaHandlerRoutingBuilder', () => {
   it('should build routing config', () => {
     type ExpectedMessage<MessageValue> = Message<string, MessageValue, string, string>
 
     // Given
-    const builder = new KafkaHandlerRoutingBuilder<TopicsConfig>()
+    const builder = new KafkaHandlerRoutingBuilder<TopicsConfig, ExecutionContext>()
       .addConfig(
         'all',
-        new KafkaHandlerConfig(CREATE_SCHEMA, (message, requestContext) => {
+        new KafkaHandlerConfig(CREATE_SCHEMA, (message, executionContext, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof CREATE_SCHEMA>>>()
+          expectTypeOf(executionContext).toEqualTypeOf<ExecutionContext>()
           expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
       .addConfig(
         'all',
-        new KafkaHandlerConfig(UPDATE_SCHEMA, (message, requestContext) => {
+        new KafkaHandlerConfig(UPDATE_SCHEMA, (message, executionContext, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof UPDATE_SCHEMA>>>()
+          expectTypeOf(executionContext).toEqualTypeOf<ExecutionContext>()
           expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
       .addConfig(
         'empty',
-        new KafkaHandlerConfig(EMPTY_SCHEMA, (message, requestContext) => {
+        new KafkaHandlerConfig(EMPTY_SCHEMA, (message, executionContext, requestContext) => {
           expectTypeOf(message).toEqualTypeOf<ExpectedMessage<z.infer<typeof EMPTY_SCHEMA>>>()
+          expectTypeOf(executionContext).toEqualTypeOf<ExecutionContext>()
           expectTypeOf(requestContext).toEqualTypeOf<RequestContext>()
         }),
       )
diff --git a/packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.ts b/packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.ts
@@ -8,23 +8,27 @@ import type { KafkaHandlerConfig } from './KafkaHandlerConfig.ts'
 
 export type KafkaHandlerRouting<
   TopicsConfig extends TopicConfig[],
+  ExecutionContext,
   MessageValue extends SupportedMessageValues<TopicsConfig> = SupportedMessageValues<TopicsConfig>,
-> = Record<string, KafkaHandlerConfig<MessageValue>[]>
+> = Record<string, KafkaHandlerConfig<MessageValue, ExecutionContext>[]>
 
-export class KafkaHandlerRoutingBuilder<TopicsConfig extends TopicConfig[]> {
-  private readonly configs: KafkaHandlerRouting<TopicsConfig> = {}
+export class KafkaHandlerRoutingBuilder<
+  const TopicsConfig extends TopicConfig[],
+  ExecutionContext,
+> {
+  private readonly configs: KafkaHandlerRouting<TopicsConfig, ExecutionContext> = {}
 
   addConfig<
     Topic extends SupportedTopics<TopicsConfig>,
     MessageValue extends SupportedMessageValuesForTopic<TopicsConfig, Topic>,
-  >(topic: Topic, config: KafkaHandlerConfig<MessageValue>): this {
+  >(topic: Topic, config: KafkaHandlerConfig<MessageValue, ExecutionContext>): this {
     this.configs[topic] ??= []
     this.configs[topic].push(config)
 
     return this
   }
 
-  build(): KafkaHandlerRouting<TopicsConfig> {
+  build(): KafkaHandlerRouting<TopicsConfig, ExecutionContext> {
     return this.configs
   }
 }
diff --git a/packages/kafka/package.json b/packages/kafka/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@message-queue-toolkit/kafka",
-    "version": "0.1.3",
+    "version": "0.2.0",
     "engines": {
         "node": ">= 22.14.0"
     },
diff --git a/packages/kafka/test/consumer/PermissionConsumer.spec.ts b/packages/kafka/test/consumer/PermissionConsumer.spec.ts
@@ -324,10 +324,10 @@ describe('PermissionConsumer', () => {
       consumer = new PermissionConsumer(testContext.cradle, {
         handlers: {
           'permission-general': [
-            new KafkaHandlerConfig(PERMISSION_ADDED_SCHEMA, (message, requestContext) => {
+            new KafkaHandlerConfig(PERMISSION_ADDED_SCHEMA, (message, _, requestContext) => {
               handlerCalls.push({ messageValue: message.value, requestContext })
             }),
-            new KafkaHandlerConfig(PERMISSION_SCHEMA, (message, requestContext) => {
+            new KafkaHandlerConfig(PERMISSION_SCHEMA, (message, _, requestContext) => {
               handlerCalls.push({ messageValue: message.value, requestContext })
             }),
           ],
@@ -373,7 +373,7 @@ describe('PermissionConsumer', () => {
         headerRequestIdField,
         handlers: {
           'permission-general': [
-            new KafkaHandlerConfig(PERMISSION_ADDED_SCHEMA, (message, requestContext) => {
+            new KafkaHandlerConfig(PERMISSION_ADDED_SCHEMA, (message, _, requestContext) => {
               handlerCalls.push({ messageValue: message.value, requestContext })
             }),
           ],
diff --git a/packages/kafka/test/consumer/PermissionConsumer.ts b/packages/kafka/test/consumer/PermissionConsumer.ts

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@message-queue-toolkit/kafka",`
`3`		`- "version": "0.1.3",`
	`3`	`+ "version": "0.2.0",`
`4`	`4`	`"engines": {`
`5`	`5`	`"node": ">= 22.14.0"`
`6`	`6`	`},`