feat: add Kafka Avro reader

pyproject.toml

@@ -35,6 +35,7 @@ dev = [
     "pytest-mock>=3.14.0",
     "confluent-kafka>=2.5.3",
     "fastavro>=1.9.7",
+    "types-confluent-kafka>=1.2.2",
 ]
 [tool.commitizen]
 version = "1.0.0"

src/sparkle/reader/kafka_reader.py

@@ -1,3 +1,4 @@
+from typing import Any
 from pyspark.sql import SparkSession, DataFrame
 from sparkle.config import Config
 from sparkle.reader.schema_registry import SchemaRegistry
@@ -25,7 +26,7 @@ def __init__(
         schema_registry: SchemaRegistry,
         use_avro: bool = True,
         schema_version: str = "latest",
-        kafka_spark_options: dict[str, str] = {},
+        kafka_spark_options: dict[str, Any] = {},
     ):
         """Initializes KafkaReader with configuration, Spark session, topic, and schema registry.
 

tests/unit/reader/test_kafka_reader.py

@@ -2,7 +2,7 @@
 from collections.abc import Generator
 import pytest
 from pyspark.sql import SparkSession, DataFrame
-from confluent_kafka import Producer, KafkaException
+from confluent_kafka import Producer
 from confluent_kafka.admin import AdminClient, NewTopic
 from confluent_kafka.schema_registry import SchemaRegistryClient, Schema
 from confluent_kafka.schema_registry.avro import AvroSerializer
@@ -29,14 +29,9 @@ def kafka_setup() -> Generator[str, None, None]:
     """
     admin_client = AdminClient({"bootstrap.servers": KAFKA_BROKER_URL})
 
-    # Create topic if it does not exist
-    try:
-        admin_client.create_topics(
-            [NewTopic(TEST_TOPIC, num_partitions=1, replication_factor=1)]
-        )
-    except KafkaException as e:
-        if e.args[0].code() != KafkaException.TOPIC_ALREADY_EXISTS:
-            raise e
+    admin_client.create_topics(
+        [NewTopic(TEST_TOPIC, num_partitions=1, replication_factor=1)]
+    )
 
     yield TEST_TOPIC
 
@@ -88,7 +83,7 @@ def avro_serializer(schema_registry_client: SchemaRegistryClient) -> AvroSeriali
     schema = Schema(schema_str, schema_type="AVRO")
     schema_registry_client.register_schema(f"{TEST_TOPIC}-value", schema)
 
-    return AvroSerializer(schema_registry_client, schema_str, lambda obj, ctx: obj)
+    return AvroSerializer(schema_registry_client, schema_str)
 
 
 @pytest.fixture

tests/unit/writer/test_kafka_writer.py

@@ -6,6 +6,7 @@
 from pyspark.sql.functions import floor, rand
 from pyspark.sql import DataFrame
 from pyspark.sql import SparkSession
+import time
 
 
 @pytest.fixture
@@ -23,7 +24,7 @@ def kafka_config() -> dict[str, Any]:
             "kafka.security.protocol": "PLAINTEXT",
         },
         "checkpoint_location": "./tmp/checkpoint",
-        "kafka_topic": "test_topic",
+        "kafka_topic": "test-kafka-writer-topic",
         "output_mode": "append",
         "unique_identifier_column_name": "id",
         "trigger_once": True,
@@ -108,6 +109,8 @@ def test_kafka_stream_publisher_write(
     except Exception as e:
         pytest.fail(f"KafkaStreamPublisher write failed with exception: {e}")
 
+    # Wait to make sure commit file is created
+    time.sleep(5)
     checkpoint_dir = kafka_config["checkpoint_location"]
     commit_file_path = os.path.join(checkpoint_dir, "commits", "0")
     assert os.path.exists(commit_file_path), "Commit file does not exist"