Add pass-through Kafka producer config via --kafka-config

- Add --kafka-config for inline JSON config
- Add --kafka-config-file for JSON file config
- KafkaLoader passes extra config to KafkaProducer
- Supports all kafka-python options (compression, SSL, SASL, etc.)

Author: incrypto32

apps/kafka_streaming_loader.py

@@ -2,6 +2,7 @@
 """Stream data to Kafka with resume watermark support."""
 
 import argparse
+import json
 import logging
 import os
 import time
@@ -82,6 +83,7 @@ def main(
     auth_token: str = None,
     max_retries: int = 5,
     retry_delay: float = 1.0,
+    kafka_config: dict = None,
 ):
     def connect():
         return Client(amp_server, auth=auth, auth_token=auth_token)
@@ -98,15 +100,14 @@ def connect():
     else:
         label_config = None
 
-    client.configure_connection(
-        'kafka',
-        'kafka',
-        {
-            'bootstrap_servers': kafka_brokers,
-            'client_id': 'amp-kafka-loader',
-            'state': {'enabled': True, 'storage': 'lmdb', 'data_dir': state_dir},
-        },
-    )
+    connection_config = {
+        'bootstrap_servers': kafka_brokers,
+        'client_id': 'amp-kafka-loader',
+        'state': {'enabled': True, 'storage': 'lmdb', 'data_dir': state_dir},
+    }
+    if kafka_config:
+        connection_config.update(kafka_config)
+    client.configure_connection('kafka', 'kafka', connection_config)
 
     with open(query_file) as f:
         query = f.read()
@@ -154,13 +155,32 @@ def stream_batches():
     parser.add_argument('--auth-token', help='Explicit auth token (works independently, does not require --auth)')
     parser.add_argument('--max-retries', type=int, default=5, help='Max retries for connection failures (default: 5)')
     parser.add_argument('--retry-delay', type=float, default=1.0, help='Initial retry delay in seconds (default: 1.0)')
+    parser.add_argument(
+        '--kafka-config',
+        type=str,
+        help='Extra Kafka producer config as JSON. Uses kafka-python naming (underscores). '
+        'Example: \'{"compression_type": "lz4", "linger_ms": 5}\'. '
+        'See: https://kafka-python.readthedocs.io/en/master/apidoc/KafkaProducer.html',
+    )
+    parser.add_argument(
+        '--kafka-config-file',
+        type=Path,
+        help='Path to JSON file with extra Kafka producer config',
+    )
     parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'])
     args = parser.parse_args()
 
     logging.basicConfig(level=logging.WARNING, format='%(asctime)s [%(name)s] %(levelname)s: %(message)s')
     log_level = getattr(logging, args.log_level) if args.log_level else logging.INFO
     logging.getLogger('amp').setLevel(log_level)
 
+    kafka_config = {}
+    if args.kafka_config_file:
+        kafka_config = json.loads(args.kafka_config_file.read_text())
+        logger.info(f'Loaded Kafka config from {args.kafka_config_file}')
+    if args.kafka_config:
+        kafka_config.update(json.loads(args.kafka_config))
+
     try:
         main(
             amp_server=args.amp_server,
@@ -176,6 +196,7 @@ def stream_batches():
             auth_token=args.auth_token,
             max_retries=args.max_retries,
             retry_delay=args.retry_delay,
+            kafka_config=kafka_config or None,
         )
     except KeyboardInterrupt:
         logger.info('Stopped by user')

src/amp/loaders/implementations/kafka_loader.py

@@ -16,8 +16,9 @@ class KafkaConfig:
     client_id: str = 'amp-kafka-loader'
     key_field: Optional[str] = 'id'
 
-    def __post_init__(self):
-        pass
+
+KAFKA_CONFIG_FIELDS = {'bootstrap_servers', 'client_id', 'key_field'}
+RESERVED_CONFIG_FIELDS = {'resilience', 'state', 'checkpoint', 'idempotency'}
 
 
 class KafkaLoader(DataLoader[KafkaConfig]):
@@ -26,6 +27,9 @@ class KafkaLoader(DataLoader[KafkaConfig]):
     SUPPORTS_TRANSACTIONS = True
 
     def __init__(self, config: Dict[str, Any], label_manager=None) -> None:
+        self._extra_producer_config = {
+            k: v for k, v in config.items() if k not in KAFKA_CONFIG_FIELDS and k not in RESERVED_CONFIG_FIELDS
+        }
         super().__init__(config, label_manager)
         self._producer = None
 
@@ -34,12 +38,16 @@ def _get_required_config_fields(self) -> list[str]:
 
     def connect(self) -> None:
         try:
-            self._producer = KafkaProducer(
-                bootstrap_servers=self.config.bootstrap_servers,
-                client_id=self.config.client_id,
-                value_serializer=lambda x: json.dumps(x, default=str).encode('utf-8'),
-                transactional_id=f'{self.config.client_id}-txn',
-            )
+            producer_config = {
+                'bootstrap_servers': self.config.bootstrap_servers,
+                'client_id': self.config.client_id,
+                'value_serializer': lambda x: json.dumps(x, default=str).encode('utf-8'),
+                'transactional_id': f'{self.config.client_id}-txn',
+                **self._extra_producer_config,
+            }
+            if self._extra_producer_config:
+                self.logger.info(f'Extra Kafka config: {list(self._extra_producer_config.keys())}')
+            self._producer = KafkaProducer(**producer_config)
 
             self._producer.init_transactions()