Examples

Admin: create a topic, then produce

from pyrula.kafka import (
    KafkaAdminClient, KafkaConfig, KafkaProducer,
    TopicSpec,
)

cfg = KafkaConfig(bootstrap_servers="localhost:9092")
admin = KafkaAdminClient(cfg)

result = admin.create_topics([
    TopicSpec("orders", num_partitions=4, replication_factor=1),
])
if result.is_err():
    raise RuntimeError(result.error)

with KafkaProducer(cfg) as producer:
    receipt = producer.produce_many_json("orders", values=[{"id": 1}, {"id": 2}])
    receipt.get().unwrap_or_raise()

Consumer group with manual commit

from pyrula.kafka import KafkaConsumerConfig, KafkaConsumer

cfg = KafkaConsumerConfig(
    bootstrap_servers="localhost:9092",
    topics=["orders"],
    group_id="order-processor",
    enable_auto_commit=False,
    auto_offset_reset="earliest",
    max_poll_records=200,
)

with KafkaConsumer(cfg) as consumer:
    while True:
        result = consumer.poll_batch(max_records=200, timeout_ms=1000)
        if result.is_err():
            print("poll error:", result.error)
            continue
        batch = result.value
        if batch.count == 0:
            continue
        for record in batch.records.to_list():
            process(record)           # your logic here
        consumer.commit(batch).unwrap_or_raise()  # advances committed offsets on the broker

Transactional consume-transform-produce

Read, transform, write, and commit the input offsets in the same transaction. On failure, abort and rewind so the batch is delivered again.

from pyrula.kafka import (
    KafkaConfig, KafkaConsumerConfig,
    KafkaProducer, KafkaConsumer,
    KafkaRecord,
)

producer_cfg = KafkaConfig(
    bootstrap_servers="localhost:9092",
    transactional_id="etl-worker-1",
)
consumer_cfg = KafkaConsumerConfig(
    bootstrap_servers="localhost:9092",
    topics=["raw-events"],
    group_id="etl-group",
    enable_auto_commit=False,
    isolation_level="read_committed",
)

producer = KafkaProducer(producer_cfg)
producer.init_transactions().unwrap_or_raise()

with KafkaConsumer(consumer_cfg) as consumer:
    while True:
        result = consumer.poll_batch(max_records=500, timeout_ms=1000)
        if result.is_err() or result.value.count == 0:
            continue
        batch = result.value

        enriched = [
            KafkaRecord("enriched-events", enrich(r.value), key=r.key)
            for r in batch.records.to_list()
        ]

        producer.begin_transaction().unwrap_or_raise()
        # Returns a receipt at once; the commit waits for every record.
        producer.produce_many_records_txn(enriched)
        group = consumer.consumer_group_metadata()  # Nothing until the consumer joins
        producer.send_offsets_to_transaction(batch.offsets, group).unwrap_or_raise()
        committed = producer.commit_transaction()
        if committed.is_err() and committed.error.txn_requires_abort():
            producer.abort_transaction().unwrap_or_raise()
            consumer.rewind(batch).unwrap_or_raise()  # reprocess this batch
        else:
            committed.unwrap_or_raise()

producer.close()

Avro schema from a dataclass, produce and consume

import dataclasses
from pyrula import AvroSerde
from pyrula.kafka import (
    KafkaConfig, KafkaConsumerConfig,
    KafkaProducer, KafkaConsumer,
    avro_schema_json,
)

@dataclasses.dataclass
class Order:
    id: int
    total: float

SCHEMA = avro_schema_json(Order)
serde = AvroSerde(SCHEMA)

config = KafkaConfig(bootstrap_servers="localhost:9092")

with KafkaProducer(config) as producer:
    receipt = producer.produce_many_avro(
        "orders-avro",
        values=[{"id": 1, "total": 42.5}, {"id": 2, "total": 9.99}],
        serde=serde,
    )
    receipt.get().unwrap_or_raise()

consumer_cfg = KafkaConsumerConfig(
    bootstrap_servers="localhost:9092",
    topics=["orders-avro"],
    group_id="avro-reader",
    auto_offset_reset="earliest",
)

with KafkaConsumer(consumer_cfg) as consumer:
    result = consumer.poll_batch(max_records=10, timeout_ms=2000)
    if result.is_ok():
        for record in result.value.records.to_list():
            order = serde.deserialize(record.value)
            print(order.get_or_else(None))
        consumer.commit(result.value).unwrap_or_raise()