Add --reorg-topic option for separate reorg message topic

Allows reorg messages to be sent to a different Kafka topic than data
messages. Required for Snowflake Kafka connector which needs strict
schema per topic. Defaults to same topic if not specified.

Author: incrypto32

apps/kafka_streaming_loader.py

@@ -84,6 +84,7 @@ def main(
     max_retries: int = 5,
     retry_delay: float = 1.0,
     kafka_config: dict = None,
+    reorg_topic: str = None,
 ):
     def connect():
         return Client(amp_server, auth=auth, auth_token=auth_token)
@@ -105,6 +106,8 @@ def connect():
         'client_id': 'amp-kafka-loader',
         'state': {'enabled': True, 'storage': 'lmdb', 'data_dir': state_dir},
     }
+    if reorg_topic:
+        connection_config['reorg_topic'] = reorg_topic
     if kafka_config:
         connection_config.update(kafka_config)
     client.configure_connection('kafka', 'kafka', connection_config)
@@ -146,6 +149,7 @@ def stream_batches():
     parser.add_argument('--amp-server', default=os.getenv('AMP_SERVER_URL', 'grpc://127.0.0.1:1602'))
     parser.add_argument('--kafka-brokers', default='localhost:9092')
     parser.add_argument('--topic', required=True)
+    parser.add_argument('--reorg-topic', help='Separate topic for reorg messages (default: same as --topic)')
     parser.add_argument('--query-file', required=True)
     parser.add_argument(
         '--raw-dataset', required=True, help='Dataset name for the raw dataset of the chain (e.g., anvil, eth_firehose)'
@@ -200,6 +204,7 @@ def stream_batches():
             max_retries=args.max_retries,
             retry_delay=args.retry_delay,
             kafka_config=kafka_config or None,
+            reorg_topic=args.reorg_topic,
         )
     except KeyboardInterrupt:
         logger.info('Stopped by user')

src/amp/loaders/implementations/kafka_loader.py

@@ -15,9 +15,10 @@ class KafkaConfig:
     bootstrap_servers: str
     client_id: str = 'amp-kafka-loader'
     key_field: Optional[str] = 'id'
+    reorg_topic: Optional[str] = None
 
 
-KAFKA_CONFIG_FIELDS = {'bootstrap_servers', 'client_id', 'key_field'}
+KAFKA_CONFIG_FIELDS = {'bootstrap_servers', 'client_id', 'key_field', 'reorg_topic'}
 RESERVED_CONFIG_FIELDS = {'resilience', 'state', 'checkpoint', 'idempotency'}
 
 
@@ -126,14 +127,14 @@ def _extract_message_key(self, row: Dict[str, Any]) -> Optional[bytes]:
 
     def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str, connection_name: str) -> None:
         """
-        Handle blockchain reorganization by sending reorg events to the same topic.
+        Handle blockchain reorganization by sending reorg events to Kafka.
 
         Reorg events are sent as special messages with _type='reorg' so consumers
         can detect and handle invalidated block ranges.
 
         Args:
             invalidation_ranges: List of block ranges to invalidate
-            table_name: The Kafka topic name to send reorg events to
+            table_name: The Kafka topic name (used if reorg_topic not configured)
             connection_name: Connection identifier (unused for Kafka, but required by base class)
         """
         if not invalidation_ranges:
@@ -143,6 +144,8 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str,
             self.logger.warning('Producer not connected, skipping reorg handling')
             return
 
+        reorg_topic = self.config.reorg_topic or table_name
+
         self._producer.begin_transaction()
         try:
             for invalidation_range in invalidation_ranges:
@@ -155,16 +158,16 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str,
                 }
 
                 self._producer.send(
-                    topic=table_name, key=f'reorg:{invalidation_range.network}'.encode('utf-8'), value=reorg_message
+                    topic=reorg_topic, key=f'reorg:{invalidation_range.network}'.encode('utf-8'), value=reorg_message
                 )
 
                 self.logger.info(
-                    f'Sent reorg event to {table_name}: '
+                    f'Sent reorg event to {reorg_topic}: '
                     f'{invalidation_range.network} blocks {invalidation_range.start}-{invalidation_range.end}'
                 )
 
             self._producer.commit_transaction()
-            self.logger.info(f'Committed {len(invalidation_ranges)} reorg events to {table_name}')
+            self.logger.info(f'Committed {len(invalidation_ranges)} reorg events to {reorg_topic}')
 
         except Exception as e:
             self._producer.abort_transaction()

tests/integration/test_kafka_loader.py

@@ -148,6 +148,52 @@ def test_handle_reorg(self, kafka_test_config):
         assert msg2.value['end_block'] == 600
         assert msg2.value['last_valid_hash'] == '0xdef456'
 
+    def test_handle_reorg_separate_topic(self, kafka_test_config):
+        config_with_reorg_topic = {
+            **kafka_test_config,
+            'reorg_topic': 'test_reorg_events',
+        }
+        loader = KafkaLoader(config_with_reorg_topic)
+        data_topic = 'test_data_topic_separate'
+
+        batch = pa.RecordBatch.from_pydict({'id': [1, 2], 'value': [100, 200]})
+        invalidation_ranges = [
+            BlockRange(network='ethereum', start=100, end=200, hash='0xabc123'),
+        ]
+
+        with loader:
+            loader.load_batch(batch, data_topic)
+            loader._handle_reorg(invalidation_ranges, data_topic, 'test_connection')
+
+        data_consumer = KafkaConsumer(
+            data_topic,
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+        data_messages = list(data_consumer)
+        data_consumer.close()
+
+        assert len(data_messages) == 2
+        assert all(msg.value['_type'] == 'data' for msg in data_messages)
+
+        reorg_consumer = KafkaConsumer(
+            'test_reorg_events',
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+        reorg_messages = list(reorg_consumer)
+        reorg_consumer.close()
+
+        assert len(reorg_messages) == 1
+        assert reorg_messages[0].value['_type'] == 'reorg'
+        assert reorg_messages[0].value['network'] == 'ethereum'
+        assert reorg_messages[0].value['start_block'] == 100
+        assert reorg_messages[0].value['end_block'] == 200
+
     def test_streaming_with_reorg(self, kafka_test_config):
         loader = KafkaLoader(kafka_test_config)
         topic_name = 'test_streaming_topic'