use lmdb stream state for kafka loader

Author: incrypto32

.gitignore

@@ -60,3 +60,4 @@ data/
 # Build artifacts
 *.tar.gz
 *.zip
+.amp_state/

apps/kafka_streaming_loader.py

@@ -57,7 +57,15 @@ def main(
     else:
         label_config = None
 
-    client.configure_connection('kafka', 'kafka', {'bootstrap_servers': kafka_brokers, 'client_id': 'amp-kafka-loader'})
+    client.configure_connection(
+        'kafka',
+        'kafka',
+        {
+            'bootstrap_servers': kafka_brokers,
+            'client_id': 'amp-kafka-loader',
+            'state': {'enabled': True, 'storage': 'lmdb'},
+        },
+    )
 
     with open(query_file) as f:
         query = f.read()

src/amp/loaders/implementations/kafka_loader.py

@@ -5,6 +5,7 @@
 import pyarrow as pa
 from kafka import KafkaProducer
 
+from ...streaming.lmdb_state import LMDBStreamStateStore
 from ...streaming.types import BlockRange
 from ..base import DataLoader, LoadMode
 
@@ -47,6 +48,10 @@ def connect(self) -> None:
             self.logger.info(f'Connected to Kafka at {self.config.bootstrap_servers}')
             self.logger.info(f'Client ID: {self.config.client_id}')
 
+            if self.state_enabled and self.state_storage == 'lmdb':
+                self.state_store = LMDBStreamStateStore(connection_name=self.config.client_id)
+                self.logger.info(f'Initialized LMDB state store at {self.state_store.data_dir}')
+
             self._is_connected = True
 
         except Exception as e:
@@ -58,6 +63,9 @@ def disconnect(self) -> None:
             self._producer.close()
             self._producer = None
 
+        if isinstance(self.state_store, LMDBStreamStateStore):
+            self.state_store.close()
+
         self._is_connected = False
         self.logger.info('Disconnected from Kafka')