initialize the kafka loader with some tests

Author: incrypto32

.gitignore

@@ -0,0 +1,88 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# Project specific
+*.log
+.env.local
+.env.*.local

pyproject.toml

@@ -65,14 +65,19 @@ 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>=3.5.0",  # Snowflake
-    "lmdb>=1.4.0",                  # LMDB
+    "pydantic>=2.0,<2.12", # PyIceberg 0.10.0 compatibility
+    "snowflake-connector-python>=3.5.0", # Snowflake
+    "lmdb>=1.4.0", # LMDB
+    "kafka-python>=2.2.15",
 ]
 
 test = [

src/amp/loaders/implementations/__init__.py

@@ -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')

src/amp/loaders/implementations/kafka_loader.py

@@ -0,0 +1,66 @@
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import pyarrow as pa
+from kafka import KafkaProducer
+
+from ...streaming.types import BlockRange
+from ..base import DataLoader, LoadMode
+
+
+@dataclass
+class KafkaConfig:
+    bootstrap_servers: str
+    client_id: str = 'amp-kafka-loader'
+    key_field: Optional[str] = 'id'
+
+    def __post_init__(self):
+        pass
+
+
+class KafkaLoader(DataLoader[KafkaConfig]):
+    SUPPORTED_MODES = {LoadMode.APPEND}
+    REQUIRES_SCHEMA_MATCH = False
+    SUPPORTS_TRANSACTIONS = False
+
+    def __init__(self, config: Dict[str, Any]) -> None:
+        super().__init__(config)
+        self._producer = None
+
+    def _get_required_config_fields(self) -> list[str]:
+        return ['bootstrap_servers']
+
+    def connect(self) -> None:
+        try:
+            self._producer = KafkaProducer(
+                bootstrap_servers=self.config.bootstrap_servers,
+                client_id=self.config.client_id,
+                value_serializer=lambda x: json.dumps(x, default=str).encode('utf-8'),
+            )
+
+            metadata = self._producer.bootstrap_connected()
+            self.logger.info(f'Connection status: {metadata}')
+            self.logger.info(f'Connected to Kafka at {self.config.bootstrap_servers}')
+            self.logger.info(f'Client ID: {self.config.client_id}')
+
+            self._is_connected = True
+
+        except Exception as e:
+            self.logger.error(f'Failed to connect to Kafka: {e}')
+            raise
+
+    def disconnect(self) -> None:
+        if self._producer:
+            self._producer.flush()
+            self._producer.close()
+            self._producer = None
+
+        self._is_connected = False
+        self.logger.info('Disconnected from Kafka')
+
+    def _load_batch_impl(self, batch: pa.RecordBatch, table_name: str, **kwargs) -> int:
+        pass
+
+    def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str) -> None:
+        pass

tests/conftest.py

@@ -34,6 +34,7 @@
     try:
         from testcontainers.postgres import PostgresContainer
         from testcontainers.redis import RedisContainer
+        from testcontainers.kafka import KafkaContainer
 
         TESTCONTAINERS_AVAILABLE = True
     except ImportError:
@@ -201,6 +202,46 @@ def redis_test_config(request):
         return request.getfixturevalue('redis_config')
 
 
+@pytest.fixture(scope='session')
+def kafka_container():
+    """Kafka container for integration tests"""
+    if not TESTCONTAINERS_AVAILABLE:
+        pytest.skip('Testcontainers not available')
+
+    container = KafkaContainer(image='confluentinc/cp-kafka:7.6.0')
+    container.start()
+
+    time.sleep(10)
+
+    yield container
+
+    container.stop()
+
+
+@pytest.fixture(scope='session')
+def kafka_config():
+    """Kafka configuration from environment or defaults"""
+    return {
+        'bootstrap_servers': os.getenv('KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092'),
+        'client_id': 'amp-test-client',
+    }
+
+
+@pytest.fixture(scope='session')
+def kafka_test_config(request):
+    """Kafka configuration from testcontainer or environment"""
+    if TESTCONTAINERS_AVAILABLE and USE_TESTCONTAINERS:
+        kafka_container = request.getfixturevalue('kafka_container')
+        bootstrap_servers = kafka_container.get_bootstrap_server()
+
+        return {
+            'bootstrap_servers': bootstrap_servers,
+            'client_id': 'amp-test-client',
+        }
+    else:
+        return request.getfixturevalue('kafka_config')
+
+
 @pytest.fixture(scope='session')
 def delta_test_env():
     """Create Delta Lake test environment for the session"""
@@ -408,6 +449,7 @@ def pytest_configure(config):
     config.addinivalue_line('markers', 'performance: Performance and benchmark tests')
     config.addinivalue_line('markers', 'postgresql: Tests requiring PostgreSQL')
     config.addinivalue_line('markers', 'redis: Tests requiring Redis')
+    config.addinivalue_line('markers', 'kafka: Tests requiring Apache Kafka')
     config.addinivalue_line('markers', 'delta_lake: Tests requiring Delta Lake')
     config.addinivalue_line('markers', 'iceberg: Tests requiring Apache Iceberg')
     config.addinivalue_line('markers', 'snowflake: Tests requiring Snowflake')

tests/integration/test_kafka_loader.py

@@ -0,0 +1,31 @@
+import pytest
+
+try:
+    from src.amp.loaders.implementations.kafka_loader import KafkaLoader
+except ImportError:
+    pytest.skip('amp modules not available', allow_module_level=True)
+
+
+@pytest.mark.integration
+@pytest.mark.kafka
+class TestKafkaLoaderIntegration:
+
+    def test_loader_connection(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+
+        loader.connect()
+        assert loader._is_connected == True
+        assert loader._producer is not None
+
+        loader.disconnect()
+        assert loader._is_connected == False
+        assert loader._producer is None
+
+    def test_context_manager(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+
+        with loader:
+            assert loader._is_connected == True
+            assert loader._producer is not None
+
+        assert loader._is_connected == False