Refactor kafka streaming loader and fix Resume watermark

Author: incrypto32

apps/kafka_streaming_loader.py

@@ -1,187 +1,124 @@
 #!/usr/bin/env python3
-"""
-Kafka streaming loader with label joining.
-Continuously loads ERC20 transfers to Kafka with token metadata.
-"""
+"""Stream data to Kafka with resume watermark support."""
 
 import argparse
-import json
+import logging
 import os
-import time
 from pathlib import Path
 
 from amp.client import Client
 from amp.loaders.types import LabelJoinConfig
-from kafka import KafkaConsumer
-
-
-def consume_messages(kafka_brokers: str, topic: str, max_messages: int = 10):
-    """Consume and print messages from Kafka topic for testing."""
-    print(f'\n{"=" * 60}')
-    print('Consuming messages from Kafka')
-    print(f'{"=" * 60}\n')
-    print(f'Topic: {topic}')
-    print(f'Brokers: {kafka_brokers}')
-    print(f'Max messages: {max_messages}\n')
-
-    consumer = KafkaConsumer(
-        topic,
-        bootstrap_servers=kafka_brokers,
-        auto_offset_reset='earliest',
-        value_deserializer=lambda x: json.loads(x.decode('utf-8')),
-        consumer_timeout_ms=5000,
-        group_id='kafka-streaming-loader-consumer',
-        enable_auto_commit=True,
-    )
+from amp.streaming import BlockRange, ResumeWatermark
+
+
+def get_block_hash(client: Client, raw_dataset: str, block_num: int) -> str:
+    """Get block hash from dataset.blocks table."""
+    query = f'SELECT hash FROM {raw_dataset}.blocks WHERE block_num = {block_num} LIMIT 1'
+    result = client.get_sql(query, read_all=True)
+    hash_val = result.to_pydict()['hash'][0]
+    return '0x' + hash_val.hex() if isinstance(hash_val, bytes) else hash_val
 
-    def format_address(addr):
-        """Convert binary address to hex string."""
-        if addr is None:
-            return None
-        if isinstance(addr, str):
-            return addr
-        if isinstance(addr, bytes):
-            return '0x' + addr.hex()
-        return addr
-
-    msg_count = 0
-    for message in consumer:
-        msg_count += 1
-        data = message.value
-        print(f'Message {msg_count}:')
-        print(f'  block_num: {data.get("block_num")}')
-        print(f'  token_address: {format_address(data.get("token_address"))}')
-        print(f'  symbol: {data.get("symbol")}')
-        print(f'  name: {data.get("name")}')
-        print(f'  decimals: {data.get("decimals")}')
-        print(f'  value: {data.get("value")}')
-        print(f'  from_address: {format_address(data.get("from_address"))}')
-        print(f'  to_address: {format_address(data.get("to_address"))}')
-        print()
-
-        if msg_count >= max_messages:
-            break
-
-    consumer.close()
-    print(f'Consumed {msg_count} messages from Kafka topic "{topic}"')
+
+def get_latest_block(client: Client, raw_dataset: str) -> int:
+    """Get latest block number from dataset.blocks table."""
+    query = f'SELECT block_num FROM {raw_dataset}.blocks ORDER BY block_num DESC LIMIT 1'
+    result = client.get_sql(query, read_all=True)
+    return result.to_pydict()['block_num'][0]
+
+
+def create_watermark(client: Client, raw_dataset: str, network: str, start_block: int) -> ResumeWatermark:
+    """Create a resume watermark for the given start block."""
+    watermark_block = start_block - 1
+    watermark_hash = get_block_hash(client, raw_dataset, watermark_block)
+    return ResumeWatermark(
+        ranges=[BlockRange(network=network, start=watermark_block, end=watermark_block, hash=watermark_hash)]
+    )
 
 
 def main(
+    amp_server: str,
     kafka_brokers: str,
     topic: str,
-    label_csv: str,
-    amp_server: str,
     query_file: str,
-    consume_mode: bool = False,
-    consume_max: int = 10,
+    raw_dataset: str,
+    network: str,
+    start_block: int = None,
+    label_csv: str = None,
 ):
-    if consume_mode:
-        consume_messages(kafka_brokers, topic, consume_max)
-        return
-
-    print(f'Connecting to Amp server: {amp_server}')
     client = Client(amp_server)
+    print(f'Connected to {amp_server}')
 
-    label_path = Path(label_csv)
-    if not label_path.exists():
-        raise FileNotFoundError(f'Label CSV not found: {label_csv}')
+    if label_csv and Path(label_csv).exists():
+        client.configure_label('tokens', label_csv)
+        print(f'Loaded {len(client.label_manager.get_label("tokens"))} labels from {label_csv}')
+        label_config = LabelJoinConfig(
+            label_name='tokens', label_key_column='token_address', stream_key_column='token_address'
+        )
+    else:
+        label_config = None
 
-    client.configure_label('tokens', str(label_path))
-    print(f'Loaded {len(client.label_manager.get_label("tokens"))} tokens from {label_csv}')
+    client.configure_connection('kafka', 'kafka', {'bootstrap_servers': kafka_brokers, 'client_id': 'amp-kafka-loader'})
 
-    kafka_config = {
-        'bootstrap_servers': kafka_brokers,
-        'client_id': 'amp-kafka-loader',
-    }
-    client.configure_connection('kafka', 'kafka', kafka_config)
+    with open(query_file) as f:
+        query = f.read()
 
-    query_path = Path(query_file)
-    if not query_path.exists():
-        raise FileNotFoundError(f'Query file not found: {query_file}')
+    if start_block is None:
+        start_block = get_latest_block(client, raw_dataset) - 10
 
-    with open(query_path) as f:
-        query = f.read()
+    print(f'Starting from block {start_block}')
 
-    label_config = LabelJoinConfig(
-        label_name='tokens',
-        label_key_column='token_address',
-        stream_key_column='token_address',
-    )
+    resume_watermark = create_watermark(client, raw_dataset, network, start_block) if start_block > 0 else None
+    if resume_watermark:
+        print(f'Watermark: {resume_watermark.to_json()}')
 
-    print(f'Starting Kafka streaming loader')
-    print(f'Kafka brokers: {kafka_brokers}')
-    print(f'Topic: {topic}')
-    print('Press Ctrl+C to stop\n')
+    print(f'Streaming to Kafka: {kafka_brokers} -> {topic}\n')
 
-    total_rows = 0
     batch_count = 0
-
     for result in client.sql(query).load(
-        connection='kafka',
-        destination=topic,
-        stream=True,
-        label_config=label_config,
+        'kafka', topic, stream=True, label_config=label_config, resume_watermark=resume_watermark
     ):
         if result.success:
-            total_rows += result.rows_loaded
             batch_count += 1
-            print(f'Batch {batch_count}: {result.rows_loaded} rows in {result.duration:.2f}s (total: {total_rows})')
+            if batch_count == 1 and result.metadata:
+                print(f'First batch: {result.metadata.get("block_ranges")}\n')
+            print(f'Batch {batch_count}: {result.rows_loaded} rows in {result.duration:.2f}s')
         else:
             print(f'Error: {result.error}')
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Stream ERC20 transfers to Kafka with token labels')
-    parser.add_argument(
-        '--kafka-brokers',
-        default=os.getenv('KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092'),
-        help='Kafka bootstrap servers (default: localhost:9092 or KAFKA_BOOTSTRAP_SERVERS env var)',
-    )
-    parser.add_argument('--topic', default='erc20_transfers', help='Kafka topic name (default: erc20_transfers)')
+    parser = argparse.ArgumentParser(description='Stream data to Kafka with resume watermark')
+    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('--query-file', required=True)
     parser.add_argument(
-        '--label-csv',
-        default='data/eth_mainnet_token_metadata.csv',
-        help='Path to token metadata CSV (default: data/eth_mainnet_token_metadata.csv)',
+        '--raw-dataset', required=True, help='Dataset name for the raw dataset of the chain (e.g., anvil, eth_firehose)'
     )
-    parser.add_argument(
-        '--amp-server',
-        default=os.getenv('AMP_SERVER_URL', 'grpc://34.27.238.174:80'),
-        help='Amp server URL (default: grpc://34.27.238.174:80 or AMP_SERVER_URL env var)',
-    )
-    parser.add_argument(
-        '--query-file',
-        default='apps/queries/erc20_transfers.sql',
-        help='Path to SQL query file (default: apps/queries/erc20_transfers.sql)',
-    )
-    parser.add_argument(
-        '--consume',
-        action='store_true',
-        help='Consume mode: read and print messages from Kafka topic (for testing)',
-    )
-    parser.add_argument(
-        '--consume-max',
-        type=int,
-        default=10,
-        help='Maximum messages to consume in consume mode (default: 10)',
-    )
-
+    parser.add_argument('--network', default='anvil')
+    parser.add_argument('--start-block', type=int, help='Start from specific block (default: latest - 10)')
+    parser.add_argument('--label-csv', help='Optional CSV for label joining')
+    parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'])
     args = parser.parse_args()
 
+    if args.log_level:
+        logging.basicConfig(
+            level=getattr(logging, args.log_level), format='%(asctime)s [%(name)s] %(levelname)s: %(message)s'
+        )
+
     try:
         main(
+            amp_server=args.amp_server,
             kafka_brokers=args.kafka_brokers,
             topic=args.topic,
-            label_csv=args.label_csv,
-            amp_server=args.amp_server,
             query_file=args.query_file,
-            consume_mode=args.consume,
-            consume_max=args.consume_max,
+            raw_dataset=args.raw_dataset,
+            network=args.network,
+            start_block=args.start_block,
+            label_csv=args.label_csv,
         )
     except KeyboardInterrupt:
-        print('\n\nInterrupted by user')
+        print('\n\nStopped by user')
     except Exception as e:
-        print(f'\n\nError: {e}')
-        import traceback
-
-        traceback.print_exc()
+        print(f'\nError: {e}')
         raise

apps/queries/anvil_logs.sql

@@ -0,0 +1,7 @@
+SELECT
+    block_num,
+    tx_hash,
+    log_index,
+    address,
+    topic0
+FROM anvil.logs

apps/queries/erc20_transfers_streaming.sql

@@ -3,7 +3,11 @@
 Test ERC20 query with label joining
 """
 
+import json
 import os
+import time
+
+from kafka import KafkaConsumer
 
 from amp.client import Client
 from amp.loaders.types import LabelJoinConfig
@@ -76,10 +80,6 @@
 print('Reading back from Kafka')
 print('=' * 60)
 
-from kafka import KafkaConsumer
-import json
-import time
-
 time.sleep(1)
 
 consumer = KafkaConsumer(
@@ -90,7 +90,7 @@
     consumer_timeout_ms=5000,
 )
 
-print(f'\nConsuming messages from topic "erc20_transfers":\n')
+print('\nConsuming messages from topic "erc20_transfers":\n')
 msg_count = 0
 for message in consumer:
     msg_count += 1