Kafka zod 4 (#317)

* Deps update

* Zod import update

* Minor error fixes

* Migration fixes

* Import fix

* More fixes

* More fixes

* Fixing zod issues

* Fixing types

* Release prepare

* Fixing tests

* Lint fixes

Author: CarlosGamero

packages/kafka/lib/AbstractKafkaConsumer.ts

@@ -19,10 +19,16 @@ import {
 import { AbstractKafkaService, type BaseKafkaOptions } from './AbstractKafkaService.ts'
 import { KafkaHandlerContainer } from './handler-container/KafkaHandlerContainer.ts'
 import type { KafkaHandlerRouting } from './handler-container/KafkaHandlerRoutingBuilder.ts'
-import type { KafkaHandler, RequestContext } from './handler-container/index.ts'
-import type { KafkaConfig, KafkaDependencies, TopicConfig } from './types.ts'
-import { ILLEGAL_GENERATION, REBALANCE_IN_PROGRESS, UNKNOWN_MEMBER_ID } from './utils/errorCodes.js'
-import { safeJsonDeserializer } from './utils/safeJsonDeserializer.js'
+import type { KafkaHandler } from './handler-container/index.ts'
+import type {
+  KafkaConfig,
+  KafkaDependencies,
+  RequestContext,
+  SupportedMessageValues,
+  TopicConfig,
+} from './types.ts'
+import { ILLEGAL_GENERATION, REBALANCE_IN_PROGRESS, UNKNOWN_MEMBER_ID } from './utils/errorCodes.ts'
+import { safeJsonDeserializer } from './utils/safeJsonDeserializer.ts'
 
 export type KafkaConsumerDependencies = KafkaDependencies &
   Pick<QueueConsumerDependencies, 'transactionObservabilityManager'>
@@ -128,7 +134,11 @@ export abstract class AbstractKafkaConsumer<
       })
     }
 
-    this.consumerStream.on('data', (message) => this.consume(message))
+    this.consumerStream.on('data', (message) =>
+      this.consume(
+        message as Message<string, SupportedMessageValues<TopicsConfig>, string, string>,
+      ),
+    )
     this.consumerStream.on('error', (error) => this.handlerError(error))
   }
 
@@ -140,7 +150,9 @@ export abstract class AbstractKafkaConsumer<
     await this.consumer.close()
   }
 
-  private async consume(message: Message<string, object, string, string>): Promise<void> {
+  private async consume(
+    message: Message<string, SupportedMessageValues<TopicsConfig>, string, string>,
+  ): Promise<void> {
     // message.value can be undefined if the message is not JSON-serializable
     if (!message.value) return this.commitMessage(message)
 
@@ -260,7 +272,9 @@ export abstract class AbstractKafkaConsumer<
     }
   }
 
-  private buildTransactionName(message: Message<string, object, string, string>) {
+  private buildTransactionName(
+    message: Message<string, SupportedMessageValues<TopicsConfig>, string, string>,
+  ) {
     const messageType = this.resolveMessageType(message.value)
 
     let name = `kafka:${this.constructor.name}:${message.topic}`

packages/kafka/lib/AbstractKafkaPublisher.ts

@@ -8,10 +8,10 @@ import {
   stringSerializer,
 } from '@platformatic/kafka'
 import { AbstractKafkaService, type BaseKafkaOptions } from './AbstractKafkaService.ts'
-import type { RequestContext } from './handler-container/index.js'
 import type {
   KafkaDependencies,
-  SupportedMessageValuesInput,
+  RequestContext,
+  SupportedMessageValuesForTopic,
   SupportedMessageValuesInputForTopic,
   SupportedTopics,
   TopicConfig,
@@ -31,10 +31,7 @@ export abstract class AbstractKafkaPublisher<
   TopicsConfig extends TopicConfig[],
 > extends AbstractKafkaService<TopicsConfig, KafkaPublisherOptions<TopicsConfig>> {
   private readonly topicsConfig: TopicsConfig
-  private readonly schemaContainers: Record<
-    string,
-    MessageSchemaContainer<SupportedMessageValuesInput<TopicsConfig>>
-  >
+  private readonly schemaContainers: Record<string, MessageSchemaContainer<object>>
 
   private readonly producer: Producer<string, object, string, string>
   private isInitiated: boolean
@@ -97,14 +94,17 @@ export abstract class AbstractKafkaPublisher<
   ): Promise<void> {
     const messageProcessingStartTimestamp = Date.now()
 
-    const schemaResult = this.schemaContainers[topic]?.resolveSchema(message)
+    const schemaResult = this.schemaContainers[topic]?.resolveSchema(message as object)
     if (!schemaResult) throw new Error(`Message schemas not found for topic: ${topic}`)
     if (schemaResult.error) throw schemaResult.error
 
     await this.init() // lazy init
 
     try {
-      const parsedMessage = schemaResult.result.parse(message)
+      const parsedMessage = schemaResult.result.parse(message) as SupportedMessageValuesForTopic<
+        TopicsConfig,
+        Topic
+      >
 
       const headers = {
         ...options?.headers,

packages/kafka/lib/AbstractKafkaService.ts

@@ -36,8 +36,11 @@ export type BaseKafkaOptions = {
   logMessages?: boolean
 } & Omit<BaseOptions, keyof KafkaConfig> // Exclude properties that are already in KafkaConfig
 
-type ProcessedMessage = MayOmit<
-  Pick<Message<string, object, string, string>, 'topic' | 'value' | 'timestamp'>,
+type ProcessedMessage<TopicsConfig extends TopicConfig[]> = MayOmit<
+  Pick<
+    Message<string, SupportedMessageValues<TopicsConfig>, string, string>,
+    'topic' | 'value' | 'timestamp'
+  >,
   'timestamp'
 >
 
@@ -76,6 +79,7 @@ export abstract class AbstractKafkaService<
 
   protected resolveMessageType(message: SupportedMessageValues<TopicsConfig>): string | undefined {
     if (!this.options.messageTypeField) return undefined
+    // @ts-expect-error
     return message[this.options.messageTypeField] as string | undefined
   }
 
@@ -89,7 +93,7 @@ export abstract class AbstractKafkaService<
   }
 
   protected handleMessageProcessed(params: {
-    message: ProcessedMessage
+    message: ProcessedMessage<TopicsConfig>
     processingResult: MessageProcessingResult
     messageProcessingStartTimestamp: number
   }) {

packages/kafka/lib/handler-container/KafkaHandlerConfig.ts

@@ -1,11 +1,6 @@
-import type { CommonLogger } from '@lokalise/node-core'
 import type { Message } from '@platformatic/kafka'
-import type { ZodSchema, ZodTypeDef } from 'zod'
-
-export interface RequestContext {
-  logger: CommonLogger
-  reqId: string
-}
+import type { ZodSchema } from 'zod/v4'
+import type { RequestContext } from '../types.js'
 
 export type KafkaHandler<MessageValue extends object, ExecutionContext> = (
   message: Message<string, MessageValue, string, string>,
@@ -14,11 +9,13 @@ export type KafkaHandler<MessageValue extends object, ExecutionContext> = (
 ) => Promise<void> | void
 
 export class KafkaHandlerConfig<MessageValue extends object, ExecutionContext> {
-  public readonly schema: ZodSchema<MessageValue, ZodTypeDef, unknown>
+  // biome-ignore lint/suspicious/noExplicitAny: Input for schema is flexible
+  public readonly schema: ZodSchema<MessageValue, any>
   public readonly handler: KafkaHandler<MessageValue, ExecutionContext>
 
   constructor(
-    schema: ZodSchema<MessageValue, ZodTypeDef, unknown>,
+    // biome-ignore lint/suspicious/noExplicitAny: Input for schema is flexible
+    schema: ZodSchema<MessageValue, any>,
     handler: KafkaHandler<MessageValue, ExecutionContext>,
   ) {
     this.schema = schema

packages/kafka/lib/handler-container/KafkaHandlerContainer.spec.ts

@@ -1,8 +1,8 @@
-import z from 'zod/v3'
+import z from 'zod/v4'
 import type { TopicConfig } from '../types.ts'
 import { KafkaHandlerConfig } from './KafkaHandlerConfig.ts'
 import { KafkaHandlerContainer } from './KafkaHandlerContainer.ts'
-import type { KafkaHandlerRouting } from './KafkaHandlerRoutingBuilder.ts'
+import { KafkaHandlerRoutingBuilder } from './KafkaHandlerRoutingBuilder.ts'
 
 const CREATE_SCHEMA = z.object({
   type: z.literal('create'),
@@ -44,32 +44,33 @@ describe('KafkaHandlerContainer', () => {
 
       // When & Then
       expect(
-        () => new KafkaHandlerContainer(topicHandlers1, 'type'),
+        () => new KafkaHandlerContainer(topicHandlers1 as any, 'type'),
       ).toThrowErrorMatchingInlineSnapshot('[Error: Duplicate handler for topic create]')
-      expect(() => new KafkaHandlerContainer(topicHandlers2)).toThrowErrorMatchingInlineSnapshot(
-        '[Error: Duplicate handler for topic empty]',
-      )
+      expect(
+        () => new KafkaHandlerContainer(topicHandlers2 as any),
+      ).toThrowErrorMatchingInlineSnapshot('[Error: Duplicate handler for topic empty]')
     })
 
     it('should resolve handler with message type', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
-        all: [
-          new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()),
-          new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()),
-          new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()),
-        ],
-        create: [new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve())],
-        empty: [new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve())],
-      }
+      const routing = new KafkaHandlerRoutingBuilder<TopicsConfig, any>()
+        .addConfig('all', new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('all', new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('all', new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()))
+        .addConfig('create', new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('empty', new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()))
 
       // When
-      const container = new KafkaHandlerContainer(topicHandlers, 'type')
+      const container = new KafkaHandlerContainer(routing.build(), 'type')
 
       // Then
-      expect(container.resolveHandler('all', { type: 'create' })?.schema).toBe(CREATE_SCHEMA)
+      expect(container.resolveHandler('all', { type: 'create', prop: 1 })?.schema).toBe(
+        CREATE_SCHEMA,
+      )
       expect(container.resolveHandler('all', { type: 'update' })?.schema).toBe(UPDATE_SCHEMA)
-      expect(container.resolveHandler('all', { type: 'non-existing' })?.schema).toBe(EMPTY_SCHEMA)
+      expect(container.resolveHandler('all', { type: 'non-existing' as any })?.schema).toBe(
+        EMPTY_SCHEMA,
+      )
       expect(container.resolveHandler('all', {})?.schema).toBe(EMPTY_SCHEMA)
 
       expect(container.resolveHandler('create', { type: 'create', prop: 1 })?.schema).toBe(
@@ -81,18 +82,19 @@ describe('KafkaHandlerContainer', () => {
       expect(container.resolveHandler('create', {} as any)?.schema).toBe(undefined)
 
       expect(container.resolveHandler('empty', {} as any)?.schema).toBe(EMPTY_SCHEMA)
-      expect(container.resolveHandler('empty', { type: 'create' })?.schema).toBe(EMPTY_SCHEMA)
+      expect(container.resolveHandler('empty', { type: 'create' } as any)?.schema).toBe(
+        EMPTY_SCHEMA,
+      )
     })
 
     it('should resolve handler without message type', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
-        create: [new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve())],
-        empty: [new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve())],
-      }
+      const routing = new KafkaHandlerRoutingBuilder<TopicsConfig, any>()
+        .addConfig('create', new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('empty', new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()))
 
       // When
-      const container = new KafkaHandlerContainer(topicHandlers)
+      const container = new KafkaHandlerContainer(routing.build())
 
       // Then
       expect(container.resolveHandler('create', { type: 'create', prop: 1 })?.schema).toBe(
@@ -104,7 +106,9 @@ describe('KafkaHandlerContainer', () => {
       expect(container.resolveHandler('create', {} as any)?.schema).toBe(CREATE_SCHEMA)
 
       expect(container.resolveHandler('empty', {} as any)?.schema).toBe(EMPTY_SCHEMA)
-      expect(container.resolveHandler('empty', { type: 'create' })?.schema).toBe(EMPTY_SCHEMA)
+      expect(container.resolveHandler('empty', { type: 'create' } as any)?.schema).toBe(
+        EMPTY_SCHEMA,
+      )
     })
   })
 
@@ -119,19 +123,15 @@ describe('KafkaHandlerContainer', () => {
 
     it('should return all topics', () => {
       // Given
-      const topicHandlers: KafkaHandlerRouting<TopicsConfig, any, any> = {
-        all: [
-          new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()),
-          new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()),
-          new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()),
-        ],
-        create: [new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve())],
-        empty: [new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve())],
-        another: [],
-      }
+      const routing = new KafkaHandlerRoutingBuilder<TopicsConfig, any>()
+        .addConfig('all', new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('all', new KafkaHandlerConfig(UPDATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('all', new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()))
+        .addConfig('create', new KafkaHandlerConfig(CREATE_SCHEMA, () => Promise.resolve()))
+        .addConfig('empty', new KafkaHandlerConfig(EMPTY_SCHEMA, () => Promise.resolve()))
 
       // When
-      const container = new KafkaHandlerContainer(topicHandlers, 'type')
+      const container = new KafkaHandlerContainer({ ...routing.build(), another: [] }, 'type')
 
       // Then
       expect(container.topics).toEqual(['all', 'create', 'empty'])

packages/kafka/lib/handler-container/KafkaHandlerContainer.ts

@@ -65,6 +65,7 @@ export class KafkaHandlerContainer<TopicsConfig extends TopicConfig[], Execution
     if (!handlers) return undefined
 
     let messageValueType: string | undefined = undefined
+    // @ts-expect-error
     if (this.messageTypeField) messageValueType = messageValue[this.messageTypeField]
 
     return messageValueType

packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.spec.ts

@@ -1,8 +1,8 @@
 import type { Message } from '@platformatic/kafka'
 import { expectTypeOf } from 'vitest'
-import z from 'zod'
-import type { TopicConfig } from '../types.ts'
-import { KafkaHandlerConfig, type RequestContext } from './KafkaHandlerConfig.ts'
+import z from 'zod/v4'
+import type { RequestContext, TopicConfig } from '../types.ts'
+import { KafkaHandlerConfig } from './KafkaHandlerConfig.ts'
 import { KafkaHandlerRoutingBuilder } from './KafkaHandlerRoutingBuilder.ts'
 
 const CREATE_SCHEMA = z.object({ type: z.literal('create') })

packages/kafka/lib/handler-container/KafkaHandlerRoutingBuilder.ts

@@ -6,11 +6,10 @@ import type {
 } from '../types.ts'
 import type { KafkaHandlerConfig } from './KafkaHandlerConfig.ts'
 
-export type KafkaHandlerRouting<
-  TopicsConfig extends TopicConfig[],
-  ExecutionContext,
-  MessageValue extends SupportedMessageValues<TopicsConfig> = SupportedMessageValues<TopicsConfig>,
-> = Record<string, KafkaHandlerConfig<MessageValue, ExecutionContext>[]>
+export type KafkaHandlerRouting<TopicsConfig extends TopicConfig[], ExecutionContext> = Record<
+  string,
+  KafkaHandlerConfig<SupportedMessageValues<TopicsConfig>, ExecutionContext>[]
+>
 
 export class KafkaHandlerRoutingBuilder<
   const TopicsConfig extends TopicConfig[],
@@ -23,7 +22,9 @@ export class KafkaHandlerRoutingBuilder<
     MessageValue extends SupportedMessageValuesForTopic<TopicsConfig, Topic>,
   >(topic: Topic, config: KafkaHandlerConfig<MessageValue, ExecutionContext>): this {
     this.configs[topic] ??= []
-    this.configs[topic].push(config)
+    this.configs[topic].push(
+      config as KafkaHandlerConfig<SupportedMessageValues<TopicsConfig>, ExecutionContext>,
+    )
 
     return this
   }

packages/kafka/lib/types.ts

@@ -1,7 +1,12 @@
+import type { CommonLogger } from '@lokalise/node-core'
 import type { QueueDependencies } from '@message-queue-toolkit/core'
 import type { ConnectionOptions } from '@platformatic/kafka'
-import type { ZodSchema } from 'zod'
-import type z from 'zod/v3'
+import type { ZodSchema, z } from 'zod/v4'
+
+export interface RequestContext {
+  logger: CommonLogger
+  reqId: string
+}
 
 export type KafkaDependencies = QueueDependencies
 
@@ -12,7 +17,7 @@ export type KafkaConfig = {
 
 export type TopicConfig<Topic extends string = string> = {
   topic: Topic
-  schemas: ZodSchema[]
+  schemas: ZodSchema<object>[]
 }
 
 export type SupportedTopics<TopicsConfig extends TopicConfig[]> = TopicsConfig[number]['topic']
@@ -34,6 +39,6 @@ type MessageSchemas<TopicsConfig extends TopicConfig[]> = TopicsConfig[number]['
 export type SupportedMessageValuesInput<TopicsConfig extends TopicConfig[]> = z.input<
   MessageSchemas<TopicsConfig>
 >
-export type SupportedMessageValues<TopicsConfig extends TopicConfig[]> = z.input<
+export type SupportedMessageValues<TopicsConfig extends TopicConfig[]> = z.output<
   MessageSchemas<TopicsConfig>
 >

packages/kafka/lib/utils/safeJsonDeserializer.spec.ts

@@ -1,5 +1,5 @@
 import { stringValueSerializer } from '@lokalise/node-core'
-import { safeJsonDeserializer } from './safeJsonDeserializer.js'
+import { safeJsonDeserializer } from './safeJsonDeserializer.ts'
 
 describe('safeJsonDeserializer', () => {
   it('should deserialize valid JSON strings', () => {