Tutorial: Kafka to S3 Pipeline¶

#!/usr/bin/env python3
"""
Kafka to S3 Pipeline

Reads user events from Kafka, transforms them, and writes to S3.
"""

import json
import os
from datetime import datetime
from typing import Iterator

import pyarrow as pa
import pyarrow.parquet as pq
from kafka import KafkaConsumer
import boto3


def get_kafka_consumer() -> KafkaConsumer:
    """Create Kafka consumer from environment."""
    return KafkaConsumer(
        os.environ["INPUT_TOPIC"],
        bootstrap_servers=os.environ["KAFKA_BOOTSTRAP_SERVERS"],
        group_id=os.environ.get("CONSUMER_GROUP", "kafka-to-s3-consumer"),
        auto_offset_reset="earliest",
        value_deserializer=lambda x: json.loads(x.decode("utf-8")),
        consumer_timeout_ms=30000,  # 30 second timeout
    )


def get_s3_client():
    """Create S3 client from environment."""
    return boto3.client(
        "s3",
        endpoint_url=os.environ.get("S3_ENDPOINT"),
        aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID", "minioadmin"),
        aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY", "minioadmin"),
    )


def transform_event(event: dict) -> dict:
    """Transform a single event."""
    return {
        "event_id": event.get("id"),
        "user_id": event.get("user_id"),
        "event_type": event.get("type"),
        "timestamp": event.get("timestamp"),
        "processed_at": datetime.utcnow().isoformat(),
        "properties": json.dumps(event.get("properties", {})),
    }


def batch_events(consumer: KafkaConsumer, batch_size: int) -> Iterator[list]:
    """Yield batches of events from Kafka."""
    batch = []
    for message in consumer:
        batch.append(transform_event(message.value))
        if len(batch) >= batch_size:
            yield batch
            batch = []
    if batch:
        yield batch


def write_to_s3(s3_client, bucket: str, prefix: str, events: list):
    """Write events to S3 as Parquet."""
    if not events:
        return

    # Create Arrow table
    table = pa.Table.from_pylist(events)

    # Generate partition path
    date = datetime.utcnow().strftime("%Y-%m-%d")
    timestamp = datetime.utcnow().strftime("%H%M%S")
    key = f"{prefix}date={date}/events_{timestamp}.parquet"

    # Write to S3
    with pa.BufferOutputStream() as buf:
        pq.write_table(table, buf, compression="snappy")
        s3_client.put_object(
            Bucket=bucket,
            Key=key,
            Body=buf.getvalue().to_pybytes(),
        )

    print(f"Wrote {len(events)} events to s3://{bucket}/{key}")


def main():
    """Main pipeline entry point."""
    print("Starting Kafka to S3 pipeline...")

    # Configuration
    batch_size = int(os.environ.get("BATCH_SIZE", "1000"))
    bucket = os.environ["OUTPUT_BUCKET"]
    prefix = os.environ.get("OUTPUT_PREFIX", "processed/events/")

    # Initialize clients
    consumer = get_kafka_consumer()
    s3_client = get_s3_client()

    # Process batches
    total_events = 0
    for batch in batch_events(consumer, batch_size):
        write_to_s3(s3_client, bucket, prefix, batch)
        total_events += len(batch)

    print(f"Pipeline complete. Processed {total_events} events.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Generate test events for the pipeline."""

import json
import time
import uuid
from kafka import KafkaProducer

producer = KafkaProducer(
    bootstrap_servers="localhost:9092",
    value_serializer=lambda x: json.dumps(x).encode("utf-8"),
)

for i in range(100):
    event = {
        "id": str(uuid.uuid4()),
        "user_id": f"user_{i % 10}",
        "type": "page_view",
        "timestamp": time.time(),
        "properties": {
            "page": f"/page/{i}",
            "referrer": "https://example.com",
        },
    }
    producer.send("user-events", event)
    print(f"Sent event {i + 1}")

producer.flush()
print("Done!")