typedkafka 0.3.1__py3-none-any.whl

typedkafka/producer.py

@@ -0,0 +1,492 @@
+"""
+Kafka Producer with comprehensive documentation and full type safety.
+
+This module provides a well-documented, type-hinted wrapper around confluent-kafka's Producer.
+"""
+
+import json
+from typing import Any, Callable, Optional
+
+try:
+    from confluent_kafka import KafkaError as ConfluentKafkaError
+    from confluent_kafka import Producer as ConfluentProducer
+except ImportError:
+    # Make confluent-kafka optional for documentation/type checking
+    ConfluentProducer = None  # type: ignore[assignment,misc]
+    ConfluentKafkaError = None  # type: ignore[assignment,misc]
+
+from typedkafka.exceptions import ProducerError, SerializationError
+
+
+class KafkaProducer:
+    """
+    A well-documented Kafka producer with full type hints.
+
+    This class wraps confluent-kafka's Producer with:
+    - Comprehensive docstrings on every method
+    - Full type hints for IDE autocomplete
+    - Better error messages
+    - Convenient methods for common operations (send_json, send_string)
+    - Context manager support for automatic cleanup
+
+    Basic Usage:
+        >>> producer = KafkaProducer({"bootstrap.servers": "localhost:9092"})
+        >>> producer.send("my-topic", b"my message", key=b"my-key")
+        >>> producer.flush()  # Wait for all messages to be delivered
+
+    With Context Manager:
+        >>> with KafkaProducer({"bootstrap.servers": "localhost:9092"}) as producer:
+        ...     producer.send("my-topic", b"message")
+        ...     # Automatic flush and cleanup on exit
+
+    JSON Messages:
+        >>> producer.send_json("events", {"user_id": 123, "action": "click"})
+
+    Attributes:
+        config: The configuration dictionary used to initialize the producer
+    """
+
+    def __init__(self, config: dict[str, Any]):
+        """
+        Initialize a Kafka producer with the given configuration.
+
+        Args:
+            config: Configuration dictionary for the producer. Common options:
+                - bootstrap.servers (str): Comma-separated list of broker addresses
+                  Example: "localhost:9092" or "broker1:9092,broker2:9092"
+                - client.id (str): An identifier for this client
+                - acks (str|int): Number of acknowledgments the producer requires
+                  "0" = no acknowledgment, "1" = leader only, "all" = all replicas
+                - compression.type (str): Compression codec ("none", "gzip", "snappy", "lz4", "zstd")
+                - max.in.flight.requests.per.connection (int): Max unacknowledged requests
+                - linger.ms (int): Time to wait before sending a batch
+                - batch.size (int): Maximum size of a message batch in bytes
+
+        Raises:
+            ProducerError: If the producer cannot be initialized with the given config
+
+        Examples:
+            >>> # Basic producer
+            >>> producer = KafkaProducer({"bootstrap.servers": "localhost:9092"})
+
+            >>> # Producer with compression
+            >>> producer = KafkaProducer({
+            ...     "bootstrap.servers": "localhost:9092",
+            ...     "compression.type": "gzip",
+            ...     "acks": "all"
+            ... })
+        """
+        if ConfluentProducer is None:
+            raise ImportError(
+                "confluent-kafka is required. Install with: pip install confluent-kafka"
+            )
+
+        self.config = config
+        try:
+            self._producer = ConfluentProducer(config)
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to initialize Kafka producer: {e}",
+                original_error=e,
+            ) from e
+
+    def send(
+        self,
+        topic: str,
+        value: bytes,
+        key: Optional[bytes] = None,
+        partition: Optional[int] = None,
+        on_delivery: Optional[Callable[[Any, Any], None]] = None,
+    ) -> None:
+        """
+        Send a message to a Kafka topic.
+
+        This method is asynchronous - it returns immediately after queuing the message.
+        Use flush() to wait for delivery confirmation.
+
+        Args:
+            topic: The topic name to send the message to
+            value: The message payload as bytes
+            key: Optional message key as bytes. Messages with the same key go to the same partition.
+            partition: Optional partition number. If None, partition is chosen by the partitioner.
+            on_delivery: Optional callback function called when delivery succeeds or fails.
+                Signature: callback(error, message)
+
+        Raises:
+            ProducerError: If the message cannot be queued (e.g., queue is full)
+
+        Examples:
+            >>> # Send a simple message
+            >>> producer.send("my-topic", b"Hello, Kafka!")
+
+            >>> # Send with a key for partitioning
+            >>> producer.send("user-events", b"event data", key=b"user-123")
+
+            >>> # Send to a specific partition
+            >>> producer.send("my-topic", b"data", partition=0)
+
+            >>> # Send with delivery callback
+            >>> def on_delivery(err, msg):
+            ...     if err:
+            ...         print(f"Delivery failed: {err}")
+            ...     else:
+            ...         print(f"Delivered to {msg.topic()} [{msg.partition()}]")
+            >>> producer.send("topic", b"data", on_delivery=on_delivery)
+        """
+        try:
+            self._producer.produce(
+                topic=topic,
+                value=value,
+                key=key,
+                partition=partition,  # type: ignore[arg-type]
+                on_delivery=on_delivery,
+            )
+            # Poll to trigger callbacks and handle backpressure
+            self._producer.poll(0)
+        except BufferError as e:
+            raise ProducerError(
+                "Message queue is full. Try calling flush() or increasing queue.buffering.max.messages",
+                original_error=e,
+            ) from e
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to send message to topic '{topic}': {e}",
+                original_error=e,
+            ) from e
+
+    def send_json(
+        self,
+        topic: str,
+        value: Any,
+        key: Optional[str] = None,
+        partition: Optional[int] = None,
+        on_delivery: Optional[Callable[[Any, Any], None]] = None,
+    ) -> None:
+        """
+        Send a JSON-serialized message to a Kafka topic.
+
+        Convenience method that automatically serializes Python objects to JSON.
+
+        Args:
+            topic: The topic name to send the message to
+            value: Any JSON-serializable Python object (dict, list, str, int, etc.)
+            key: Optional string key (will be UTF-8 encoded)
+            partition: Optional partition number
+            on_delivery: Optional callback function for delivery confirmation
+
+        Raises:
+            SerializationError: If the value cannot be serialized to JSON
+            ProducerError: If the message cannot be queued
+
+        Examples:
+            >>> # Send a dict as JSON
+            >>> producer.send_json("events", {"user_id": 123, "action": "click"})
+
+            >>> # Send with a string key
+            >>> producer.send_json("user-data", {"name": "Alice"}, key="user-123")
+
+            >>> # Send a list
+            >>> producer.send_json("numbers", [1, 2, 3, 4, 5])
+        """
+        try:
+            value_bytes = json.dumps(value).encode("utf-8")
+        except (TypeError, ValueError) as e:
+            raise SerializationError(
+                f"Failed to serialize value to JSON: {e}",
+                value=value,
+                original_error=e,
+            ) from e
+
+        key_bytes = key.encode("utf-8") if key is not None else None
+        self.send(topic, value_bytes, key=key_bytes, partition=partition, on_delivery=on_delivery)
+
+    def send_string(
+        self,
+        topic: str,
+        value: str,
+        key: Optional[str] = None,
+        partition: Optional[int] = None,
+        on_delivery: Optional[Callable[[Any, Any], None]] = None,
+    ) -> None:
+        """
+        Send a UTF-8 encoded string message to a Kafka topic.
+
+        Convenience method for sending text messages.
+
+        Args:
+            topic: The topic name to send the message to
+            value: String message to send
+            key: Optional string key
+            partition: Optional partition number
+            on_delivery: Optional callback function for delivery confirmation
+
+        Raises:
+            ProducerError: If the message cannot be queued
+
+        Examples:
+            >>> producer.send_string("logs", "Application started successfully")
+            >>> producer.send_string("user-messages", "Hello!", key="user-123")
+        """
+        value_bytes = value.encode("utf-8")
+        key_bytes = key.encode("utf-8") if key is not None else None
+        self.send(topic, value_bytes, key=key_bytes, partition=partition, on_delivery=on_delivery)
+
+    def flush(self, timeout: float = -1) -> int:
+        """
+        Wait for all messages in the queue to be delivered.
+
+        Blocks until all messages are sent or the timeout expires.
+
+        Args:
+            timeout: Maximum time to wait in seconds. Use -1 for infinite wait (default).
+                Example: 5.0 = wait up to 5 seconds
+
+        Returns:
+            Number of messages still in queue/internal Producer state. 0 means all delivered.
+
+        Raises:
+            ProducerError: If flush fails
+
+        Examples:
+            >>> # Wait for all messages to be delivered
+            >>> producer.send("topic", b"message 1")
+            >>> producer.send("topic", b"message 2")
+            >>> remaining = producer.flush()
+            >>> if remaining == 0:
+            ...     print("All messages delivered successfully")
+
+            >>> # Wait up to 5 seconds
+            >>> remaining = producer.flush(timeout=5.0)
+            >>> if remaining > 0:
+            ...     print(f"Warning: {remaining} messages not delivered after 5 seconds")
+        """
+        try:
+            return self._producer.flush(timeout=timeout)  # type: ignore[no-any-return]
+        except Exception as e:
+            raise ProducerError(f"Flush failed: {e}", original_error=e) from e
+
+    def close(self) -> None:
+        """
+        Close the producer and release resources.
+
+        Calls flush() to ensure all queued messages are delivered before closing.
+        It's recommended to use the producer as a context manager instead of calling
+        this method directly.
+
+        Examples:
+            >>> producer = KafkaProducer({"bootstrap.servers": "localhost:9092"})
+            >>> try:
+            ...     producer.send("topic", b"message")
+            ... finally:
+            ...     producer.close()  # Ensure cleanup
+
+            >>> # Better: use context manager
+            >>> with KafkaProducer({"bootstrap.servers": "localhost:9092"}) as producer:
+            ...     producer.send("topic", b"message")
+        """
+        self.flush()
+
+    def __enter__(self) -> "KafkaProducer":
+        """
+        Enter context manager.
+
+        Returns:
+            self
+
+        Examples:
+            >>> with KafkaProducer({"bootstrap.servers": "localhost:9092"}) as producer:
+            ...     producer.send("topic", b"message")
+        """
+        return self
+
+    def send_batch(
+        self,
+        topic: str,
+        messages: list[tuple[bytes, Optional[bytes]]],
+        on_delivery: Optional[Callable[[Any, Any], None]] = None,
+    ) -> None:
+        """
+        Send a batch of messages to a Kafka topic.
+
+        Each message is a tuple of (value, key). This is more efficient than
+        calling send() repeatedly as it defers polling until after all messages
+        are queued.
+
+        Args:
+            topic: The topic name to send the messages to
+            messages: List of (value, key) tuples. Key can be None.
+            on_delivery: Optional callback for each message delivery.
+
+        Raises:
+            ProducerError: If any message cannot be queued
+
+        Examples:
+            >>> producer.send_batch("events", [
+            ...     (b"event1", b"key1"),
+            ...     (b"event2", b"key2"),
+            ...     (b"event3", None),
+            ... ])
+            >>> producer.flush()
+        """
+        for value, key in messages:
+            try:
+                self._producer.produce(
+                    topic=topic,
+                    value=value,
+                    key=key,
+                    on_delivery=on_delivery,
+                )
+            except BufferError:
+                # Flush and retry once on buffer full
+                self._producer.flush()
+                try:
+                    self._producer.produce(
+                        topic=topic,
+                        value=value,
+                        key=key,
+                        on_delivery=on_delivery,
+                    )
+                except Exception as retry_e:
+                    raise ProducerError(
+                        f"Failed to send message to topic '{topic}' after flush: {retry_e}",
+                        original_error=retry_e,
+                    ) from retry_e
+            except Exception as e:
+                raise ProducerError(
+                    f"Failed to send message to topic '{topic}': {e}",
+                    original_error=e,
+                ) from e
+        # Poll to trigger callbacks
+        self._producer.poll(0)
+
+    def init_transactions(self, timeout: float = 30.0) -> None:
+        """
+        Initialize the producer for transactions.
+
+        Must be called before any transactional methods. Requires the
+        ``transactional.id`` configuration to be set.
+
+        Args:
+            timeout: Maximum time to wait for initialization in seconds.
+
+        Raises:
+            ProducerError: If transaction initialization fails.
+
+        Examples:
+            >>> producer = KafkaProducer({
+            ...     "bootstrap.servers": "localhost:9092",
+            ...     "transactional.id": "my-txn-id",
+            ... })
+            >>> producer.init_transactions()
+        """
+        try:
+            self._producer.init_transactions(timeout)
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to initialize transactions: {e}",
+                original_error=e,
+            ) from e
+
+    def begin_transaction(self) -> None:
+        """
+        Begin a new transaction.
+
+        Raises:
+            ProducerError: If beginning the transaction fails.
+        """
+        try:
+            self._producer.begin_transaction()
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to begin transaction: {e}",
+                original_error=e,
+            ) from e
+
+    def commit_transaction(self, timeout: float = 30.0) -> None:
+        """
+        Commit the current transaction.
+
+        Args:
+            timeout: Maximum time to wait for commit in seconds.
+
+        Raises:
+            ProducerError: If committing the transaction fails.
+        """
+        try:
+            self._producer.commit_transaction(timeout)
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to commit transaction: {e}",
+                original_error=e,
+            ) from e
+
+    def abort_transaction(self, timeout: float = 30.0) -> None:
+        """
+        Abort the current transaction.
+
+        Args:
+            timeout: Maximum time to wait for abort in seconds.
+
+        Raises:
+            ProducerError: If aborting the transaction fails.
+        """
+        try:
+            self._producer.abort_transaction(timeout)
+        except Exception as e:
+            raise ProducerError(
+                f"Failed to abort transaction: {e}",
+                original_error=e,
+            ) from e
+
+    def transaction(self) -> "TransactionContext":
+        """
+        Return a context manager for transactional sends.
+
+        Automatically begins, commits, or aborts the transaction.
+
+        Returns:
+            A context manager that manages the transaction lifecycle.
+
+        Raises:
+            ProducerError: If any transaction operation fails.
+
+        Examples:
+            >>> producer = KafkaProducer({
+            ...     "bootstrap.servers": "localhost:9092",
+            ...     "transactional.id": "my-txn-id",
+            ... })
+            >>> producer.init_transactions()
+            >>> with producer.transaction():
+            ...     producer.send("topic", b"msg1")
+            ...     producer.send("topic", b"msg2")
+        """
+        return TransactionContext(self)
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """
+        Exit context manager and cleanup resources.
+
+        Automatically flushes all pending messages before exiting.
+        """
+        self.close()
+
+
+class TransactionContext:
+    """Context manager for Kafka transactions.
+
+    Begins a transaction on entry and commits on clean exit.
+    Aborts the transaction if an exception occurs.
+    """
+
+    def __init__(self, producer: KafkaProducer):
+        self._producer = producer
+
+    def __enter__(self) -> "TransactionContext":
+        self._producer.begin_transaction()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if exc_type is not None:
+            self._producer.abort_transaction()
+        else:
+            self._producer.commit_transaction()

typedkafka/retry.py

@@ -0,0 +1,154 @@
+"""
+Retry and backoff utilities for Kafka operations.
+
+Provides decorators and helpers for retrying transient Kafka failures
+with configurable backoff strategies.
+"""
+
+import functools
+import random
+import time
+from collections.abc import Sequence
+from typing import Any, Callable, Optional, TypeVar
+
+from typedkafka.exceptions import KafkaError
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def retry(
+    max_attempts: int = 3,
+    backoff_base: float = 0.5,
+    backoff_max: float = 30.0,
+    jitter: bool = True,
+    retryable_exceptions: Optional[Sequence[type[BaseException]]] = None,
+) -> Callable[[F], F]:
+    """
+    Decorator that retries a function on failure with exponential backoff.
+
+    Args:
+        max_attempts: Maximum number of attempts (including the first call).
+        backoff_base: Base delay in seconds for exponential backoff.
+        backoff_max: Maximum delay in seconds between retries.
+        jitter: If True, add random jitter to backoff delay.
+        retryable_exceptions: Exception types to retry on.
+            Defaults to ``(KafkaError,)`` if not specified.
+
+    Returns:
+        Decorated function that retries on failure.
+
+    Raises:
+        The last exception if all attempts fail.
+
+    Examples:
+        >>> from typedkafka.retry import retry
+        >>> from typedkafka.exceptions import ProducerError
+        >>>
+        >>> @retry(max_attempts=3, backoff_base=1.0)
+        ... def send_message(producer, topic, value):
+        ...     producer.send(topic, value)
+        ...     producer.flush()
+
+        >>> # Retry only specific exceptions
+        >>> @retry(max_attempts=5, retryable_exceptions=(ProducerError,))
+        ... def produce(producer, data):
+        ...     producer.send_json("events", data)
+    """
+    if retryable_exceptions is None:
+        retryable_exceptions = (KafkaError,)
+
+    retry_tuple = tuple(retryable_exceptions)
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            last_exception: Optional[BaseException] = None
+            for attempt in range(max_attempts):
+                try:
+                    return func(*args, **kwargs)
+                except retry_tuple as e:
+                    last_exception = e
+                    if attempt < max_attempts - 1:
+                        delay = min(backoff_base * (2 ** attempt), backoff_max)
+                        if jitter:
+                            delay = delay * (0.5 + random.random() * 0.5)  # noqa: S311
+                        time.sleep(delay)
+            raise last_exception  # type: ignore[misc]
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+class RetryPolicy:
+    """
+    Configurable retry policy for Kafka operations.
+
+    Provides a reusable retry configuration that can be applied to
+    multiple operations.
+
+    Examples:
+        >>> policy = RetryPolicy(max_attempts=5, backoff_base=1.0)
+        >>> result = policy.execute(lambda: producer.send("topic", b"msg"))
+    """
+
+    def __init__(
+        self,
+        max_attempts: int = 3,
+        backoff_base: float = 0.5,
+        backoff_max: float = 30.0,
+        jitter: bool = True,
+        retryable_exceptions: Optional[Sequence[type[BaseException]]] = None,
+    ):
+        """
+        Initialize a retry policy.
+
+        Args:
+            max_attempts: Maximum number of attempts.
+            backoff_base: Base delay in seconds for exponential backoff.
+            backoff_max: Maximum delay between retries.
+            jitter: If True, add random jitter to delays.
+            retryable_exceptions: Exception types to retry on.
+                Defaults to ``(KafkaError,)``.
+        """
+        self.max_attempts = max_attempts
+        self.backoff_base = backoff_base
+        self.backoff_max = backoff_max
+        self.jitter = jitter
+        self.retryable_exceptions: tuple[type[BaseException], ...] = tuple(
+            retryable_exceptions or (KafkaError,)
+        )
+
+    def execute(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        """
+        Execute a function with retry logic.
+
+        Args:
+            func: The function to execute.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            The return value of the function.
+
+        Raises:
+            The last exception if all attempts fail.
+
+        Examples:
+            >>> policy = RetryPolicy(max_attempts=3)
+            >>> policy.execute(producer.send, "topic", b"value")
+        """
+        last_exception: Optional[BaseException] = None
+        for attempt in range(self.max_attempts):
+            try:
+                return func(*args, **kwargs)
+            except self.retryable_exceptions as e:
+                last_exception = e
+                if attempt < self.max_attempts - 1:
+                    delay = min(
+                        self.backoff_base * (2 ** attempt), self.backoff_max
+                    )
+                    if self.jitter:
+                        delay = delay * (0.5 + random.random() * 0.5)  # noqa: S311
+                    time.sleep(delay)
+        raise last_exception  # type: ignore[misc]