implement a Kafka loader

Author: incrypto32

.test.env

@@ -0,0 +1,122 @@
+import marimo
+
+__generated_with = "0.17.0"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _():
+    import marimo as mo
+    from amp.client import Client
+    return Client, mo
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+    # Kafka Streaming Example
+
+    This notebook demonstrates continuous streaming from Flight SQL to Kafka with reorg detection.
+    """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""## Setup""")
+    return
+
+
+@app.cell
+def _(Client):
+    client = Client('grpc://127.0.0.1:1602')
+    return (client,)
+
+
+@app.cell
+def _(client):
+    client.configure_connection(
+        'my_kafka',
+        'kafka',
+        {
+            'bootstrap_servers': 'localhost:9092',
+            'client_id': 'amp-streaming-client',
+            'key_field': 'block_num'
+        }
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+    ## Streaming Query
+
+    This query uses `SETTINGS stream = true` to continuously stream new blocks as they arrive.
+    The loader will automatically handle blockchain reorganizations.
+    """
+    )
+    return
+
+
+@app.cell
+def _(client):
+    streaming_results = client.sql(
+        '''
+        SELECT
+            block_num,
+            log_index
+        FROM anvil.logs
+        '''
+    ).load(
+        'my_kafka',
+        'eth_logs_stream',
+        stream=True,
+        create_table=True,
+    )
+    return (streaming_results,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+    ## Monitor Stream
+
+    This cell will continuously print results as they arrive. It starts a Kafka consumer to print the results as they come in.
+    """
+    )
+    return
+
+
+@app.cell
+def _(streaming_results):
+    import threading
+    from kafka import KafkaConsumer
+    import json
+
+    def consume_kafka():
+        consumer = KafkaConsumer(
+            'eth_logs_stream',
+            bootstrap_servers='localhost:9092',
+            auto_offset_reset='latest',
+            value_deserializer=lambda m: json.loads(m.decode('utf-8'))
+        )
+        print("Kafka Consumer started")
+        for message in consumer:
+            print(f"Consumed: {message.value}")
+
+    consumer_thread = threading.Thread(target=consume_kafka, daemon=True)
+    consumer_thread.start()
+
+    print("Kafka Producer started")
+    for result in streaming_results:
+        print(f"Produced: {result}")
+    return
+
+
+if __name__ == "__main__":
+    app.run()

notebooks/kafka_streaming.py

@@ -1,6 +1,6 @@
 import marimo
 
-__generated_with = "0.14.16"
+__generated_with = "0.17.0"
 app = marimo.App(width="full")
 
 
@@ -89,15 +89,15 @@ def _(mo):
             create_table=True,
         )
     """,
-    name='_',
+    name="_"
 )
 
 
 @app.cell
 def _(psql_load_results):
     for p_result in psql_load_results:
         print(p_result)
-    return (p_result,)
+    return
 
 
 @app.cell(hide_code=True)
@@ -120,7 +120,7 @@ def _(client):
 def _(redis_load_results):
     for r_result in redis_load_results:
         print(r_result)
-    return (r_result,)
+    return
 
 
 @app.cell(hide_code=True)
@@ -149,7 +149,7 @@ def _(client):
     else:
         # Single result
         print(f'Total: {result.rows_loaded} rows')
-    return batch_result, result
+    return (batch_result,)
 
 
 @app.cell
@@ -291,7 +291,7 @@ def _(lmdb_load_result):
 def _(batch_result, lmdb_load_result):
     for lmdb_batch_result in lmdb_load_result:
             print(f'Batch: {batch_result.rows_loaded} rows')
-    return (lmdb_batch_result,)
+    return
 
 
 @app.cell
@@ -325,7 +325,7 @@ def _(env):
        myList = [ key for key, _ in txn.cursor() ]
        print(myList)
        print(len(myList))
-    return myList, txn
+    return
 
 
 @app.cell
@@ -340,7 +340,74 @@ def _(env, pa):
             batch = reader.read_next_batch()
 
             print(batch)
-    return batch, key, open_txn, reader, value
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""# Kafka""")
+    return
+
+
+@app.cell
+def _(client):
+    # Configure Kafka connection
+    client.configure_connection(
+        'my_kafka',
+        'kafka',
+        {
+            'bootstrap_servers': 'localhost:9092',
+            'client_id': 'amp-test-client',
+            'key_field': 'id'
+        }
+    )
+    return
+
+
+@app.cell
+def _(client):
+    # Load data to Kafka topic
+    kafka_load_results = client.sql('select * from eth_firehose.logs limit 100').load(
+        'my_kafka',
+        'test_logs',
+        create_table=True,
+    )
+    return (kafka_load_results,)
+
+
+@app.cell
+def _(kafka_load_results):
+    # Check results
+    for k_result in kafka_load_results:
+        print(f'Kafka batch: {k_result.rows_loaded} rows loaded, duration: {k_result.duration:.2f}s')
+    return (k_result,)
+
+
+@app.cell
+def _():
+    from kafka import KafkaConsumer
+    import json
+
+    consumer = KafkaConsumer(
+        'test_logs',
+        bootstrap_servers='localhost:9092',
+        auto_offset_reset='earliest',
+        consumer_timeout_ms=3000,
+        value_deserializer=lambda m: json.loads(m.decode('utf-8'))
+    )
+
+    messages = list(consumer)
+    consumer.close()
+
+    print(f"Total messages in Kafka: {len(messages)}")
+    print(f"\nFirst message:")
+    if messages:
+        msg = messages[0].value
+        print(f"  Block: {msg.get('block_num')}")
+        print(f"  Timestamp: {msg.get('timestamp')}")
+        print(f"  Address: {msg.get('address')}")
+
+    return
 
 
 @app.cell

notebooks/test_loaders.py

@@ -62,15 +62,20 @@ lmdb = [
     "lmdb>=1.4.0",
 ]
 
+kafka = [
+    "kafka-python>=2.2.15",
+]
+
 all_loaders = [
-    "psycopg2-binary>=2.9.0",       # PostgreSQL
-    "redis>=4.5.0",                 # Redis
-    "deltalake>=1.0.2",             # Delta Lake (consistent version)
+    "psycopg2-binary>=2.9.0", # PostgreSQL
+    "redis>=4.5.0", # Redis
+    "deltalake>=1.0.2", # Delta Lake (consistent version)
     "pyiceberg[sql-sqlite]>=0.10.0", # Apache Iceberg
     "pydantic>=2.0,<2.12",          # PyIceberg 0.10.0 compatibility
     "snowflake-connector-python>=4.0.0",  # Snowflake
     "snowpipe-streaming>=1.0.0",        # Snowpipe Streaming API
     "lmdb>=1.4.0",                  # LMDB
+    "kafka-python>=2.2.15",
 ]
 
 test = [

pyproject.toml

@@ -21,7 +21,7 @@
 
 try:
     from .iceberg_loader import IcebergLoader
-except ImportError:
+except Exception:
     IcebergLoader = None
 
 try:
@@ -34,11 +34,10 @@
 except ImportError:
     LMDBLoader = None
 
-# Add any other loaders here
-# try:
-#     from .snowflake_loader import SnowflakeLoader
-# except ImportError:
-#     SnowflakeLoader = None
+try:
+    from .kafka_loader import KafkaLoader
+except ImportError:
+    KafkaLoader = None
 
 __all__ = []
 
@@ -55,3 +54,5 @@
     __all__.append('SnowflakeLoader')
 if LMDBLoader:
     __all__.append('LMDBLoader')
+if KafkaLoader:
+    __all__.append('KafkaLoader')