implement a simple kafka loader

Author: incrypto32

.test.env

@@ -1,6 +1,6 @@
 import marimo
 
-__generated_with = "0.14.16"
+__generated_with = "0.17.0"
 app = marimo.App(width="full")
 
 
@@ -20,7 +20,7 @@ def _():
 
 @app.cell
 def _(Client):
-    client = Client('grpc://127.0.0.1')
+    client = Client('grpc://34.122.177.97:80')
     client.configure_connection('my_pg', 'postgresql', {'host': 'localhost', 'database': 'loaders_testing', 'user': 'username', 'password': 'pass', 'port': '5432'})
     client.configure_connection('my_redis', 'redis', {'host': 'localhost', 'port': 6379, 'password': 'mypassword'})
     return (client,)
@@ -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

@@ -228,7 +228,13 @@ def _load_table(
         except Exception as e:
             self.logger.error(f'Failed to load table: {e}')
             return LoadResult(
-                rows_loaded=0, duration=0.0, table_name=table_name, loader_type=loader, success=False, error=str(e)
+                rows_loaded=0,
+                duration=0.0,
+                ops_per_second=0.0,
+                table_name=table_name,
+                loader_type=loader,
+                success=False,
+                error=str(e),
             )
 
     def _load_stream(
@@ -248,7 +254,13 @@ def _load_stream(
         except Exception as e:
             self.logger.error(f'Failed to load stream: {e}')
             yield LoadResult(
-                rows_loaded=0, duration=0.0, table_name=table_name, loader_type=loader, success=False, error=str(e)
+                rows_loaded=0,
+                duration=0.0,
+                ops_per_second=0.0,
+                table_name=table_name,
+                loader_type=loader,
+                success=False,
+                error=str(e),
             )
 
     def query_and_load_streaming(

src/amp/client.py

@@ -258,7 +258,7 @@ def _validate_schema_compatibility(self, iceberg_table: IcebergTable, arrow_sche
             # Evolution mode: evolve schema to accommodate new fields
             self._evolve_schema_if_needed(iceberg_table, iceberg_schema, arrow_schema)
 
-    def _validate_schema_strict(self, iceberg_schema: IcebergSchema, arrow_schema: pa.Schema) -> None:
+    def _validate_schema_strict(self, iceberg_schema: 'IcebergSchema', arrow_schema: pa.Schema) -> None:
         """Validate schema compatibility in strict mode (no evolution)"""
         iceberg_field_names = {field.name for field in iceberg_schema.fields}
         arrow_field_names = {field.name for field in arrow_schema}
@@ -279,7 +279,7 @@ def _validate_schema_strict(self, iceberg_schema: IcebergSchema, arrow_schema: p
         self.logger.debug('Schema validation passed in strict mode')
 
     def _evolve_schema_if_needed(
-        self, iceberg_table: IcebergTable, iceberg_schema: IcebergSchema, arrow_schema: pa.Schema
+        self, iceberg_table: IcebergTable, iceberg_schema: 'IcebergSchema', arrow_schema: pa.Schema
     ) -> None:
         """Evolve the Iceberg table schema to accommodate new Arrow schema fields"""
         try:

src/amp/loaders/implementations/iceberg_loader.py

@@ -60,7 +60,37 @@ def disconnect(self) -> None:
         self.logger.info('Disconnected from Kafka')
 
     def _load_batch_impl(self, batch: pa.RecordBatch, table_name: str, **kwargs) -> int:
-        pass
+        if not self._producer:
+            raise RuntimeError('Producer not connected. Call connect() first.')
+
+        data_dict = batch.to_pydict()
+        num_rows = batch.num_rows
+
+        if num_rows == 0:
+            return 0
+
+        for i in range(num_rows):
+            row = {field: values[i] for field, values in data_dict.items()}
+
+            key = self._extract_message_key(row)
+
+            self._producer.send(topic=table_name, key=key, value=row)
+
+        self._producer.flush()
+
+        self.logger.debug(f'Sent {num_rows} messages to topic {table_name}')
+
+        return num_rows
+
+    def _extract_message_key(self, row: Dict[str, Any]) -> Optional[bytes]:
+        if not self.config.key_field or self.config.key_field not in row:
+            return None
+
+        key_value = row[self.config.key_field]
+        if key_value is None:
+            return None
+
+        return str(key_value).encode('utf-8')
 
     def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str) -> None:
         pass

src/amp/loaders/implementations/kafka_loader.py

@@ -1,3 +1,4 @@
+import pyarrow as pa
 import pytest
 
 try:
@@ -9,7 +10,6 @@
 @pytest.mark.integration
 @pytest.mark.kafka
 class TestKafkaLoaderIntegration:
-
     def test_loader_connection(self, kafka_test_config):
         loader = KafkaLoader(kafka_test_config)
 
@@ -29,3 +29,16 @@ def test_context_manager(self, kafka_test_config):
             assert loader._producer is not None
 
         assert loader._is_connected == False
+
+    def test_load_batch(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+
+        batch = pa.RecordBatch.from_pydict(
+            {'id': [1, 2, 3], 'name': ['alice', 'bob', 'charlie'], 'value': [100, 200, 300]}
+        )
+
+        with loader:
+            result = loader.load_batch(batch, 'test_topic')
+
+        assert result.success == True
+        assert result.rows_loaded == 3