Refactor KafkaLoader for robustness

Author: incrypto32

apps/kafka_streaming_loader_guide.md

@@ -188,11 +188,42 @@ uv run python apps/kafka_consumer.py anvil_logs localhost:9092 my-group
 
 ## Docker Usage
 
+### Build the loader image
+
 ```bash
-# Build image
 docker build -f Dockerfile.kafka -t amp-kafka .
+```
+
+### Quick demo: local Kafka with Docker
+
+This section runs a single-broker Kafka in Docker for quick testing. For production, point `--kafka-brokers` at your real Kafka cluster and skip this section.
+
+Start a single-broker Kafka using `confluentinc/cp-kafka`:
+
+```bash
+docker network create kafka-net
+
+docker run -d --name kafka --network kafka-net -p 9092:9092 \
+  -e KAFKA_NODE_ID=1 \
+  -e KAFKA_PROCESS_ROLES=broker,controller \
+  -e KAFKA_CONTROLLER_QUORUM_VOTERS=1@kafka:9093 \
+  -e KAFKA_LISTENERS=PLAINTEXT://0.0.0.0:29092,CONTROLLER://0.0.0.0:9093,EXTERNAL://0.0.0.0:9092 \
+  -e KAFKA_ADVERTISED_LISTENERS=PLAINTEXT://kafka:29092,EXTERNAL://localhost:9092 \
+  -e KAFKA_LISTENER_SECURITY_PROTOCOL_MAP=PLAINTEXT:PLAINTEXT,CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT \
+  -e KAFKA_CONTROLLER_LISTENER_NAMES=CONTROLLER \
+  -e KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR=1 \
+  -e KAFKA_TRANSACTION_STATE_LOG_MIN_ISR=1 \
+  -e CLUSTER_ID=MkU3OEVBNTcwNTJENDM2Qk \
+  confluentinc/cp-kafka:latest
+```
+
+This configures two listeners:
+- `kafka:29092` — for containers on the `kafka-net` network (use this from the loader)
+- `localhost:9092` — for host access (use this from `uv run` or the consumer script)
 
-# Run loader (with auth via env var)
+### Run the loader
+
+```bash
 docker run -d \
   --name amp-kafka-loader \
   --network kafka-net \
@@ -201,7 +232,7 @@ docker run -d \
   -v $(pwd)/.amp_state:/data/state \
   amp-kafka \
   --amp-server 'grpc+tls://gateway.amp.staging.thegraph.com:443' \
-  --kafka-brokers kafka:9092 \
+  --kafka-brokers kafka:29092 \
   --topic erc20_transfers \
   --query-file /data/queries/erc20_transfers_activity.sql \
   --raw-dataset 'edgeandnode/ethereum_mainnet' \
@@ -214,7 +245,11 @@ docker run -d \
 docker logs -f amp-kafka-loader
 ```
 
-Note: Ensure your Kafka container is on the same Docker network (`kafka-net`) and advertises the correct listener (`kafka:9092`).
+### Consume messages from host
+
+```bash
+uv run python apps/kafka_consumer.py erc20_transfers localhost:9092
+```
 
 ## Getting Help
 

src/amp/loaders/implementations/kafka_loader.py

@@ -1,5 +1,5 @@
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from typing import Any, Dict, List, Optional
 
 import pyarrow as pa
@@ -18,7 +18,7 @@ class KafkaConfig:
     reorg_topic: Optional[str] = None
 
 
-KAFKA_CONFIG_FIELDS = {'bootstrap_servers', 'client_id', 'key_field', 'reorg_topic'}
+KAFKA_CONFIG_FIELDS = {f.name for f in fields(KafkaConfig)}
 RESERVED_CONFIG_FIELDS = {'resilience', 'state', 'checkpoint', 'idempotency'}
 
 
@@ -40,11 +40,11 @@ def _get_required_config_fields(self) -> list[str]:
     def connect(self) -> None:
         try:
             producer_config = {
+                **self._extra_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())}')
@@ -67,6 +67,9 @@ def connect(self) -> None:
             self._is_connected = True
 
         except Exception as e:
+            if self._producer:
+                self._producer.close()
+                self._producer = None
             self.logger.error(f'Failed to connect to Kafka: {e}')
             raise
 
@@ -135,7 +138,7 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str,
         Args:
             invalidation_ranges: List of block ranges to invalidate
             table_name: The Kafka topic name (used if reorg_topic not configured)
-            connection_name: Connection identifier (unused for Kafka, but required by base class)
+            connection_name: Connection identifier (required by base class interface)
         """
         if not invalidation_ranges:
             return

src/amp/streaming/types.py

@@ -180,11 +180,3 @@ def to_json(self) -> str:
 
             data[r.network] = {'number': r.end, 'hash': r.hash}
         return json.dumps(data)
-
-    # TODO: ResumeWatermark.from_json appears to be unused. Remove?
-    @classmethod
-    def from_json(cls, json_str: str) -> 'ResumeWatermark':
-        """Deserialize from JSON string"""
-        data = json.loads(json_str)
-        ranges = [BlockRange.from_dict(r) for r in data['ranges']]
-        return cls(ranges=ranges, timestamp=data.get('timestamp'), sequence=data.get('sequence'))

tests/unit/test_crash_recovery.py

@@ -8,7 +8,6 @@
 
 import pytest
 
-from src.amp.loaders.base import LoadResult
 from src.amp.streaming.types import BlockRange, ResumeWatermark
 from tests.fixtures.mock_clients import MockDataLoader
 

tests/unit/test_streaming_types.py

@@ -369,29 +369,7 @@ def test_to_json_minimal_data(self):
         assert data['ethereum']['number'] == 200
         assert data['ethereum']['hash'] == '0xabc123'
 
-    # TODO: ResumeWatermark.from_json appears to be unused. Remove?
-    def test_from_json_full_data(self):
-        """Test deserializing watermark with all fields"""
-        json_str = json.dumps(
-            {
-                'ranges': [
-                    {'network': 'ethereum', 'start': 100, 'end': 200},
-                    {'network': 'polygon', 'start': 50, 'end': 150},
-                ],
-                'timestamp': '2024-01-01T00:00:00Z',
-                'sequence': 42,
-            }
-        )
-
-        watermark = ResumeWatermark.from_json(json_str)
-
-        assert len(watermark.ranges) == 2
-        assert watermark.ranges[0].network == 'ethereum'
-        assert watermark.timestamp == '2024-01-01T00:00:00Z'
-        assert watermark.sequence == 42
-
-    def test_round_trip_serialization(self):
-        """Test that to_json produces server format (not reversible via from_json)"""
+    def test_to_json_server_format(self):
         watermark = ResumeWatermark(
             ranges=[
                 BlockRange(network='ethereum', start=100, end=200, hash='0xabc123'),