From 90c3098b3a928d0f2e3b0dd522623bc1c296352c Mon Sep 17 00:00:00 2001 From: Gunther Xing Date: Mon, 6 Jul 2026 14:40:15 +0800 Subject: [PATCH 1/2] feat(python): add user header and origin timestamp support Add user headers to SendMessage and ReceiveMessage, expose origin timestamp, and fix Docker test infrastructure: - explicit typed header and plain Python header support - type transformation and validation across Rust and Python - user header usage examples - minor fix: Docker test config file --- core/sdk/src/prelude.rs | 6 +- examples/python/README.md | 14 + examples/python/message-headers/consumer.py | 170 ++++ examples/python/message-headers/producer.py | 237 +++++ foreign/python/Dockerfile.test | 2 +- foreign/python/apache_iggy.pyi | 400 +++++++- foreign/python/docker-compose.test.yml | 11 +- foreign/python/src/lib.rs | 5 + foreign/python/src/receive_message.rs | 22 + foreign/python/src/send_message.rs | 45 +- foreign/python/src/user_headers.rs | 857 ++++++++++++++++++ foreign/python/tests/test_consumer_group.py | 98 ++ .../python/tests/test_message_operations.py | 276 +++++- 13 files changed, 2121 insertions(+), 22 deletions(-) create mode 100644 examples/python/message-headers/consumer.py create mode 100644 examples/python/message-headers/producer.py create mode 100644 foreign/python/src/user_headers.rs diff --git a/core/sdk/src/prelude.rs b/core/sdk/src/prelude.rs index ec01c19200..f228232efc 100644 --- a/core/sdk/src/prelude.rs +++ b/core/sdk/src/prelude.rs @@ -51,9 +51,9 @@ pub use iggy_common::{ Aes256GcmEncryptor, Args, ArgsOptional, AutoLogin, CacheMetrics, CacheMetricsKey, ClientError, ClientInfoDetails, ClusterMetadata, ClusterNode, ClusterNodeRole, ClusterNodeStatus, CompressionAlgorithm, Consumer, ConsumerGroup, ConsumerGroupDetails, ConsumerGroupMember, - ConsumerKind, EncryptorKind, GlobalPermissions, HeaderKey, HeaderKind, HeaderValue, - HttpClientConfig, HttpClientConfigBuilder, HttpMethod, IdKind, Identifier, IdentityInfo, - IggyByteSize, IggyDuration, IggyError, IggyExpiry, IggyIndexView, IggyMessage, + ConsumerKind, EncryptorKind, GlobalPermissions, HeaderField, HeaderKey, HeaderKind, + HeaderValue, HttpClientConfig, HttpClientConfigBuilder, HttpMethod, IdKind, Identifier, + IdentityInfo, IggyByteSize, IggyDuration, IggyError, IggyExpiry, IggyIndexView, IggyMessage, IggyMessageHeader, IggyMessageHeaderView, IggyMessageView, IggyMessageViewIterator, IggyTimestamp, MaxTopicSize, Partition, Partitioner, Partitioning, Permissions, PersonalAccessTokenExpiry, PollMessages, PolledMessages, PollingKind, PollingStrategy, diff --git a/examples/python/README.md b/examples/python/README.md index 9bf943b75c..22fda46f25 100644 --- a/examples/python/README.md +++ b/examples/python/README.md @@ -71,6 +71,20 @@ python basic/consumer.py Demonstrates fundamental client connection, authentication, batch message sending, and polling with support for TCP/QUIC/HTTP protocols. +### Message Headers + +Shows how to attach and read Python SDK user headers with `str`, `bytes`, `bool`, `int`, and `float` values: + +```bash +# Using uv +uv run message-headers/producer.py +uv run message-headers/consumer.py + +# Without using uv +python message-headers/producer.py +python message-headers/consumer.py +``` + ## TLS Examples To test with a TLS-enabled server, start the server with TLS configured (see main README), then run: diff --git a/examples/python/message-headers/consumer.py b/examples/python/message-headers/consumer.py new file mode 100644 index 0000000000..2f96ab84cb --- /dev/null +++ b/examples/python/message-headers/consumer.py @@ -0,0 +1,170 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import argparse +import asyncio +import json +from collections.abc import Mapping +from typing import Any, NamedTuple + +from apache_iggy import ( + HeaderKey, + HeaderValue, + IggyClient, + PollingStrategy, + ReceiveMessage, + UserHeaders, +) +from loguru import logger + +STREAM_NAME = "message-headers-stream" +TOPIC_NAME = "orders" +PARTITION_ID = 0 +BATCHES_LIMIT = 5 +MESSAGES_PER_BATCH = 10 + +ORDER_CREATED_TYPE = "OrderCreated" +ORDER_CONFIRMED_TYPE = "OrderConfirmed" +ORDER_REJECTED_TYPE = "OrderRejected" + + +class ArgNamespace(NamedTuple): + connection_string: str + + +def parse_args() -> ArgNamespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "connection_string", + help=( + "Connection string for Iggy client, e.g. " + "'iggy+tcp://iggy:iggy@127.0.0.1:8090'" + ), + default="iggy+tcp://iggy:iggy@127.0.0.1:8090", + nargs="?", + type=str, + ) + return ArgNamespace(**vars(parser.parse_args())) + + +async def main(): + args: ArgNamespace = parse_args() + client = IggyClient.from_connection_string(args.connection_string) + logger.info("Connecting to Iggy") + await client.connect() + logger.info("Connected") + await consume_messages(client) + + +async def consume_messages(client: IggyClient): + interval = 0.5 + logger.info( + f"Messages will be consumed from stream: {STREAM_NAME}, " + f"topic: {TOPIC_NAME}, partition: {PARTITION_ID} " + f"with interval {interval * 1000} ms." + ) + consumed_batches = 0 + + while consumed_batches < BATCHES_LIMIT: + try: + logger.debug("Polling for messages...") + polled_messages = await client.poll_messages( + stream=STREAM_NAME, + topic=TOPIC_NAME, + partition_id=PARTITION_ID, + polling_strategy=PollingStrategy.Next(), + count=MESSAGES_PER_BATCH, + auto_commit=True, + ) + if not polled_messages: + logger.info("No messages found in current poll") + await asyncio.sleep(interval) + continue + + for message in polled_messages: + handle_message(message) + + consumed_batches += 1 + logger.info(f"Consumed {len(polled_messages)} message(s).") + await asyncio.sleep(interval) + except Exception as error: + logger.exception(f"Exception occurred while consuming messages: {error}") + break + + logger.info(f"Consumed {consumed_batches} batches of messages, exiting.") + + +def handle_message(message: ReceiveMessage): + payload = json.loads(message.payload().decode("utf-8")) + # `user_headers()` returns the explicitly typed `UserHeaders` mapping + # (a dict subclass) or None when the message carries no headers. + headers = message.user_headers() + message_type = get_message_type(headers) + + logger.info( + f"Handling message at offset {message.offset()} " + f"with origin timestamp {message.origin_timestamp()}." + ) + if headers is not None: + logger.info(f"Headers: {format_headers(headers)}") + # Opt into the convenient plain form. + logger.info(f"Plain headers: {format_headers(headers.to_plain())}") + + if message_type == ORDER_CREATED_TYPE: + handle_order_created(payload) + elif message_type == ORDER_CONFIRMED_TYPE: + handle_order_confirmed(payload) + elif message_type == ORDER_REJECTED_TYPE: + handle_order_rejected(payload) + else: + logger.warning(f"Received unknown message type: {message_type}") + + +def handle_order_created(order_created: Mapping[str, Any]): + logger.info(f"Order Created: {order_created}") + + +def handle_order_confirmed(order_confirmed: Mapping[str, Any]): + logger.info(f"Order Confirmed: {order_confirmed}") + + +def handle_order_rejected(order_rejected: Mapping[str, Any]): + logger.info(f"Order Rejected: {order_rejected}") + + +def get_message_type(headers: UserHeaders | None) -> str | None: + if headers is None: + return None + message_type = headers.get(HeaderKey.String("message-type")) + if isinstance(message_type, HeaderValue.String): + return message_type.value + return None + + +def format_headers(headers: Mapping[Any, Any]) -> dict[str, str]: + formatted: dict[str, str] = {} + for key, value in headers.items(): + formatted_key = repr(key) if isinstance(key, HeaderKey) else key + if isinstance(value, bytes): + formatted[formatted_key] = f"bytes({value.hex()})" + else: + formatted[formatted_key] = f"{value!r} ({type(value).__name__})" + return formatted + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/message-headers/producer.py b/examples/python/message-headers/producer.py new file mode 100644 index 0000000000..7fb0b2b5e5 --- /dev/null +++ b/examples/python/message-headers/producer.py @@ -0,0 +1,237 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import argparse +import asyncio +import json +import secrets +import time +from collections.abc import Mapping +from typing import NamedTuple + +from apache_iggy import HeaderKey, HeaderValue, IggyClient, StreamDetails, TopicDetails +from apache_iggy import SendMessage as Message +from loguru import logger + +STREAM_NAME = "message-headers-stream" +TOPIC_NAME = "orders" +PARTITION_ID = 0 +BATCHES_LIMIT = 5 +MESSAGES_PER_BATCH = 10 + +ORDER_CREATED_TYPE = "OrderCreated" +ORDER_CONFIRMED_TYPE = "OrderConfirmed" +ORDER_REJECTED_TYPE = "OrderRejected" + +PlainHeaderValue = str | bytes | bool | int | float +PlainHeaders = dict[str, PlainHeaderValue] +TypedHeaders = dict[HeaderKey, HeaderValue] + + +class ArgNamespace(NamedTuple): + connection_string: str + + +class SerializedMessage(NamedTuple): + message_type: str + payload: str + headers: PlainHeaders | TypedHeaders + + +class MessagesGenerator: + def __init__(self): + self.order_id = 0 + + def generate(self) -> SerializedMessage: + self.order_id += 1 + message_type = self.order_id % 3 + + if message_type == 0: + payload = { + "orderId": f"order-{self.order_id}", + "customerId": f"customer-{secrets.randbelow(100)}", + "amount": secrets.randbelow(10000) + 1, + } + return self._serialize(ORDER_CREATED_TYPE, payload) + + if message_type == 1: + payload = { + "orderId": f"order-{self.order_id // 3}", + "timestamp": int(time.time() * 1000), + } + return self._serialize(ORDER_CONFIRMED_TYPE, payload) + + payload = { + "orderId": f"order-{self.order_id // 3}", + "reason": "Insufficient balance", + } + return self._serialize(ORDER_REJECTED_TYPE, payload) + + def _serialize( + self, message_type: str, payload: Mapping[str, object] + ) -> SerializedMessage: + encoded_payload = json.dumps(payload) + if self.order_id % 5 == 0: + # Add typed headers similar to the Rust core + typed_headers: TypedHeaders = { + HeaderKey.String("message-type"): HeaderValue.String(message_type), + HeaderKey.String("content-type"): HeaderValue.String( + "application/json" + ), + HeaderKey.String("schema-version"): HeaderValue.UnsignedInt16(1), + HeaderKey.String("created-at-ms"): HeaderValue.UnsignedInt64( + int(time.time() * 1000) + ), + HeaderKey.String("retryable"): HeaderValue.Bool( + message_type == ORDER_REJECTED_TYPE + ), + HeaderKey.String("priority-score"): HeaderValue.Float32( + (self.order_id % 100) / 100 + ), + HeaderKey.String("trace-bin"): HeaderValue.Raw( + f"trace-{self.order_id}".encode() + ), + HeaderKey.UnsignedInt32(self.order_id): HeaderValue.String("order-id"), + } + return SerializedMessage(message_type, encoded_payload, typed_headers) + + # Add easy-to-use `dict[str, str | bytes | bool | int | float]` headers + # which will be translated into typed headers by the Python SDK + plain_headers: PlainHeaders = { + "message-type": message_type, + "content-type": "application/json", + "schema-version": 1, + "created-at-ms": int(time.time() * 1000), + "retryable": message_type == ORDER_REJECTED_TYPE, + "priority-score": (self.order_id % 100) / 100, + "trace-bin": f"trace-{self.order_id}".encode(), + } + return SerializedMessage(message_type, encoded_payload, plain_headers) + + +def parse_args() -> ArgNamespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "connection_string", + help=( + "Connection string for Iggy client, e.g. " + "'iggy+tcp://iggy:iggy@127.0.0.1:8090'" + ), + default="iggy+tcp://iggy:iggy@127.0.0.1:8090", + nargs="?", + type=str, + ) + return ArgNamespace(**vars(parser.parse_args())) + + +async def main(): + args: ArgNamespace = parse_args() + client = IggyClient.from_connection_string(args.connection_string) + logger.info("Connecting to Iggy") + await client.connect() + logger.info("Connected") + await init_system(client) + await produce_messages(client) + + +async def init_system(client: IggyClient): + try: + logger.info(f"Creating stream with name {STREAM_NAME}...") + stream: StreamDetails | None = await client.get_stream(STREAM_NAME) + if stream is None: + await client.create_stream(name=STREAM_NAME) + logger.info("Stream was created successfully.") + else: + logger.warning(f"Stream {stream.name} already exists with ID {stream.id}") + + except Exception as error: + logger.error(f"Error creating stream: {error}") + logger.exception(error) + + try: + logger.info(f"Creating topic {TOPIC_NAME} in stream {STREAM_NAME}") + topic: TopicDetails | None = await client.get_topic(STREAM_NAME, TOPIC_NAME) + if topic is None: + await client.create_topic( + stream=STREAM_NAME, + partitions_count=1, + name=TOPIC_NAME, + replication_factor=1, + ) + logger.info("Topic was created successfully.") + else: + logger.warning(f"Topic {topic.name} already exists with ID {topic.id}") + except Exception as error: + logger.error(f"Error creating topic {error}") + logger.exception(error) + + +async def produce_messages(client: IggyClient): + interval = 0.5 + logger.info( + f"Messages will be sent to stream: {STREAM_NAME}, " + f"topic: {TOPIC_NAME}, partition: {PARTITION_ID} " + f"with interval {interval * 1000} ms." + ) + generator = MessagesGenerator() + sent_batches = 0 + + while sent_batches < BATCHES_LIMIT: + messages: list[Message] = [] + for _ in range(MESSAGES_PER_BATCH): + serialized = generator.generate() + messages.append( + Message(serialized.payload, user_headers=serialized.headers) + ) + logger.info( + f"Prepared {serialized.message_type} with headers: " + f"{format_headers(serialized.headers)}" + ) + + try: + await client.send_messages( + stream=STREAM_NAME, + topic=TOPIC_NAME, + partitioning=PARTITION_ID, + messages=messages, + ) + sent_batches += 1 + logger.info(f"Sent {len(messages)} message(s).") + except Exception as error: + logger.error(f"Exception type: {type(error).__name__}, message: {error}") + logger.exception(error) + + await asyncio.sleep(interval) + + logger.info(f"Sent {sent_batches} batches of messages, exiting.") + + +def format_headers( + headers: Mapping[str, PlainHeaderValue] | Mapping[HeaderKey, HeaderValue], +) -> dict[str, str]: + formatted: dict[str, str] = {} + for key, value in headers.items(): + formatted_key = repr(key) if isinstance(key, HeaderKey) else key + if isinstance(value, bytes): + formatted[formatted_key] = f"bytes({value.hex()})" + else: + formatted[formatted_key] = f"{value!r} ({type(value).__name__})" + return formatted + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/foreign/python/Dockerfile.test b/foreign/python/Dockerfile.test index 777b7e98ae..ee7b17aea6 100644 --- a/foreign/python/Dockerfile.test +++ b/foreign/python/Dockerfile.test @@ -59,7 +59,7 @@ COPY foreign/python/NOTICE ./foreign/python/ COPY core/ /workspace/core/ # Copy Python SDK source (changes more frequently) -COPY foreign/python/src/ ./src/ +COPY foreign/python/src/ ./foreign/python/src/ # Install dependencies and build extension using native uv workflow WORKDIR /workspace/foreign/python diff --git a/foreign/python/apache_iggy.pyi b/foreign/python/apache_iggy.pyi index 6a9ad26ba1..8e7e21ac82 100644 --- a/foreign/python/apache_iggy.pyi +++ b/foreign/python/apache_iggy.pyi @@ -31,6 +31,8 @@ __all__ = [ "ConsumerGroup", "ConsumerGroupDetails", "ConsumerGroupMember", + "HeaderKey", + "HeaderValue", "IggyClient", "IggyConsumer", "PollingStrategy", @@ -39,6 +41,7 @@ __all__ = [ "StreamDetails", "Topic", "TopicDetails", + "UserHeaders", ] class AutoCommit: @@ -303,6 +306,356 @@ class ConsumerGroupMember: Gets the collection of partitions the consumer group member is consuming. """ +class HeaderKey: + r""" + Typed key for an Iggy user header. + + Use these constructors when the header key must preserve an explicit + wire type instead of using the common string-key dictionary form. + """ + def __hash__(self) -> builtins.int: ... + def __richcmp__(self, other: typing.Any, op: int) -> typing.Any: ... + def __repr__(self) -> builtins.str: ... + @typing.final + class Raw(HeaderKey): + r""" + Raw bytes key. The byte length must be 1..=255. + """ + + __match_args__ = ("value",) + @property + def value(self) -> bytes: ... + def __new__(cls, value: bytes) -> HeaderKey.Raw: ... + + @typing.final + class String(HeaderKey): + r""" + UTF-8 string key. The encoded byte length must be 1..=255. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.str: ... + def __new__(cls, value: builtins.str) -> HeaderKey.String: ... + + @typing.final + class Bool(HeaderKey): + r""" + Boolean key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.bool: ... + def __new__(cls, value: builtins.bool) -> HeaderKey.Bool: ... + + @typing.final + class Int8(HeaderKey): + r""" + Signed 8-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.Int8: ... + + @typing.final + class Int16(HeaderKey): + r""" + Signed 16-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.Int16: ... + + @typing.final + class Int32(HeaderKey): + r""" + Signed 32-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.Int32: ... + + @typing.final + class Int64(HeaderKey): + r""" + Signed 64-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.Int64: ... + + @typing.final + class Int128(HeaderKey): + r""" + Signed 128-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.Int128: ... + + @typing.final + class UnsignedInt8(HeaderKey): + r""" + Unsigned 8-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.UnsignedInt8: ... + + @typing.final + class UnsignedInt16(HeaderKey): + r""" + Unsigned 16-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.UnsignedInt16: ... + + @typing.final + class UnsignedInt32(HeaderKey): + r""" + Unsigned 32-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.UnsignedInt32: ... + + @typing.final + class UnsignedInt64(HeaderKey): + r""" + Unsigned 64-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.UnsignedInt64: ... + + @typing.final + class UnsignedInt128(HeaderKey): + r""" + Unsigned 128-bit integer key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderKey.UnsignedInt128: ... + + @typing.final + class Float32(HeaderKey): + r""" + 32-bit floating point key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.float: ... + def __new__(cls, value: builtins.float) -> HeaderKey.Float32: ... + + @typing.final + class Float64(HeaderKey): + r""" + 64-bit floating point key. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.float: ... + def __new__(cls, value: builtins.float) -> HeaderKey.Float64: ... + +class HeaderValue: + r""" + Typed value for an Iggy user header. + + Use these constructors when the header value must preserve an explicit + wire type instead of using the common Python scalar dictionary form. + """ + def __hash__(self) -> builtins.int: ... + def __richcmp__(self, other: typing.Any, op: int) -> typing.Any: ... + def __repr__(self) -> builtins.str: ... + @typing.final + class Raw(HeaderValue): + r""" + Raw bytes value. The byte length must be 1..=255. + """ + + __match_args__ = ("value",) + @property + def value(self) -> bytes: ... + def __new__(cls, value: bytes) -> HeaderValue.Raw: ... + + @typing.final + class String(HeaderValue): + r""" + UTF-8 string value. The encoded byte length must be 1..=255. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.str: ... + def __new__(cls, value: builtins.str) -> HeaderValue.String: ... + + @typing.final + class Bool(HeaderValue): + r""" + Boolean value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.bool: ... + def __new__(cls, value: builtins.bool) -> HeaderValue.Bool: ... + + @typing.final + class Int8(HeaderValue): + r""" + Signed 8-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.Int8: ... + + @typing.final + class Int16(HeaderValue): + r""" + Signed 16-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.Int16: ... + + @typing.final + class Int32(HeaderValue): + r""" + Signed 32-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.Int32: ... + + @typing.final + class Int64(HeaderValue): + r""" + Signed 64-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.Int64: ... + + @typing.final + class Int128(HeaderValue): + r""" + Signed 128-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.Int128: ... + + @typing.final + class UnsignedInt8(HeaderValue): + r""" + Unsigned 8-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.UnsignedInt8: ... + + @typing.final + class UnsignedInt16(HeaderValue): + r""" + Unsigned 16-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.UnsignedInt16: ... + + @typing.final + class UnsignedInt32(HeaderValue): + r""" + Unsigned 32-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.UnsignedInt32: ... + + @typing.final + class UnsignedInt64(HeaderValue): + r""" + Unsigned 64-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.UnsignedInt64: ... + + @typing.final + class UnsignedInt128(HeaderValue): + r""" + Unsigned 128-bit integer value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.int: ... + def __new__(cls, value: builtins.int) -> HeaderValue.UnsignedInt128: ... + + @typing.final + class Float32(HeaderValue): + r""" + 32-bit floating point value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.float: ... + def __new__(cls, value: builtins.float) -> HeaderValue.Float32: ... + + @typing.final + class Float64(HeaderValue): + r""" + 64-bit floating point value. + """ + + __match_args__ = ("value",) + @property + def value(self) -> builtins.float: ... + def __new__(cls, value: builtins.float) -> HeaderValue.Float64: ... + @typing.final class IggyClient: r""" @@ -691,6 +1044,11 @@ class ReceiveMessage: Retrieves the timestamp of the received message. The timestamp represents the time of the message within its topic. """ + def origin_timestamp(self) -> builtins.int: + r""" + Retrieves the origin timestamp of the received message. + The origin timestamp represents when the message was originally created. + """ def id(self) -> builtins.int: r""" Retrieves the id of the received message. @@ -710,6 +1068,10 @@ class ReceiveMessage: r""" Retrieves the partition this message belongs to. """ + def user_headers(self) -> UserHeaders | None: + r""" + Retrieves user headers attached to the received message. + """ @typing.final class SendMessage: @@ -718,7 +1080,12 @@ class SendMessage: This class wraps a Rust message meant for sending, facilitating the creation of such messages from Python and their subsequent use in Rust. """ - def __new__(cls, data: builtins.str | bytes) -> SendMessage: + def __new__( + cls, + data: builtins.str | bytes, + user_headers: dict[typing.Any, typing.Any] | None = None, + id: builtins.int | None = None, + ) -> SendMessage: r""" Constructs a new `SendMessage` instance from a string or bytes. This method allows for the creation of a `SendMessage` instance @@ -791,3 +1158,34 @@ class TopicDetails: r""" Replication factor for the topic. """ + +@typing.final +class UserHeaders(dict): + r""" + User headers dictionary returned by `ReceiveMessage.user_headers`. + + This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping + operations work) that additionally exposes `to_plain` for the convenient + scalar form. + """ + def __new__( + cls, mapping: dict[typing.Any, typing.Any] | None = None + ) -> UserHeaders: + r""" + Wraps a mapping so its entries gain the `to_plain` helper. + + Accepts a dict whose keys and values can each independently be + `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool | + int | float`). The inherited `dict` initializer copies the provided + mapping. + """ + def to_plain( + self, + ) -> dict[str | bytes | bool | int | float, str | bytes | bool | int | float]: + r""" + Converts these headers into the convenient plain dictionary form. + + Every header kind maps losslessly onto a Python scalar, so this never + loses information; it only returns an error if a stored field cannot be + decoded. + """ diff --git a/foreign/python/docker-compose.test.yml b/foreign/python/docker-compose.test.yml index 92ffec6daf..129b941bf5 100644 --- a/foreign/python/docker-compose.test.yml +++ b/foreign/python/docker-compose.test.yml @@ -15,8 +15,6 @@ # specific language governing permissions and limitations # under the License. -version: '3.8' - services: iggy-server: build: @@ -24,15 +22,22 @@ services: dockerfile: core/server/Dockerfile args: PROFILE: debug + command: ["--fresh", "--with-default-root-credentials"] container_name: iggy-server-python-test + security_opt: + - seccomp:unconfined networks: - python-test-network ports: - "3000:3000" - "8080:8080" - "8090:8090" + environment: + - IGGY_HTTP_ADDRESS=0.0.0.0:3000 + - IGGY_TCP_ADDRESS=0.0.0.0:8090 + - IGGY_QUIC_ADDRESS=0.0.0.0:8080 healthcheck: - test: [ "CMD", "curl", "-f", "http://localhost:3000/stats" ] + test: [ "CMD", "iggy", "--tcp-server-address", "127.0.0.1:8090", "ping" ] interval: 5s timeout: 5s retries: 12 diff --git a/foreign/python/src/lib.rs b/foreign/python/src/lib.rs index 8c78456062..3d22efded7 100644 --- a/foreign/python/src/lib.rs +++ b/foreign/python/src/lib.rs @@ -22,6 +22,7 @@ mod receive_message; mod send_message; mod stream; mod topic; +mod user_headers; use client::IggyClient; use consumer::{ @@ -33,6 +34,7 @@ use receive_message::{PollingStrategy, ReceiveMessage}; use send_message::SendMessage; use stream::StreamDetails; use topic::{Topic, TopicDetails}; +use user_headers::{HeaderKey, HeaderValue, UserHeaders}; /// A Python module implemented in Rust. #[pymodule] @@ -52,5 +54,8 @@ fn apache_iggy(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/foreign/python/src/receive_message.rs b/foreign/python/src/receive_message.rs index aadf0772cf..c00d748377 100644 --- a/foreign/python/src/receive_message.rs +++ b/foreign/python/src/receive_message.rs @@ -16,10 +16,13 @@ // under the License. use iggy::prelude::{IggyMessage as RustReceiveMessage, PollingStrategy as RustPollingStrategy}; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods}; +use crate::user_headers::{UserHeaders, rust_user_headers_to_py}; + /// A Python class representing a received message. /// This class wraps a Rust message, allowing for access to its payload and offset from Python. #[pyclass] @@ -50,6 +53,12 @@ impl ReceiveMessage { self.inner.header.timestamp } + /// Retrieves the origin timestamp of the received message. + /// The origin timestamp represents when the message was originally created. + pub fn origin_timestamp(&self) -> u64 { + self.inner.header.origin_timestamp + } + /// Retrieves the id of the received message. /// The id represents unique identifier of the message within its topic. pub fn id(&self) -> u128 { @@ -72,6 +81,19 @@ impl ReceiveMessage { pub fn partition_id(&self) -> u32 { self.partition_id } + + /// Retrieves user headers attached to the received message. + #[gen_stub(override_return_type(type_repr = "UserHeaders | None"))] + pub fn user_headers<'a>(&self, py: Python<'a>) -> PyResult>> { + let Some(headers) = self + .inner + .user_headers_map() + .map_err(|e| PyValueError::new_err(e.to_string()))? + else { + return Ok(None); + }; + rust_user_headers_to_py(py, headers).map(Some) + } } #[derive(Clone, Copy)] diff --git a/foreign/python/src/send_message.rs b/foreign/python/src/send_message.rs index 021125d7e6..5d96e526e9 100644 --- a/foreign/python/src/send_message.rs +++ b/foreign/python/src/send_message.rs @@ -17,12 +17,13 @@ use bytes::Bytes; use iggy::prelude::{IggyMessage as RustIggyMessage, IggyMessageHeader}; -use pyo3::{prelude::*, types::PyBytes}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyBytes}; use pyo3_stub_gen::{ derive::{gen_stub_pyclass, gen_stub_pymethods}, impl_stub_type, }; -use std::str::FromStr; + +use crate::user_headers::py_user_headers_to_rust; /// A Python class representing a message to be sent. /// This class wraps a Rust message meant for sending, facilitating @@ -61,22 +62,40 @@ impl SendMessage { /// This method allows for the creation of a `SendMessage` instance /// directly from Python using the provided string or bytes data. #[new] - pub fn new(py: Python, data: PyMessagePayload) -> PyResult { - let inner = match data { - PyMessagePayload::String(data) => RustIggyMessage::from_str(&data) - .map_err(|e| PyErr::new::(e.to_string()))?, - PyMessagePayload::Bytes(data) => { - let bytes = Bytes::from(data.extract::>(py)?); - RustIggyMessage::builder() - .payload(bytes) - .build() - .map_err(|e| PyErr::new::(e.to_string()))? - } + #[pyo3(signature = (data, user_headers=None, id=None))] + pub fn new( + py: Python, + data: PyMessagePayload, + #[gen_stub(override_type(type_repr = "dict[typing.Any, typing.Any] | None"))] + user_headers: Option<&Bound<'_, PyAny>>, + #[gen_stub(override_type(type_repr = "builtins.int | None"))] id: Option, + ) -> PyResult { + let payload = match data { + PyMessagePayload::String(data) => Bytes::from(data), + PyMessagePayload::Bytes(data) => Bytes::from(data.extract::>(py)?), }; + let user_headers = user_headers + .map(|headers| { + let headers = headers + .cast::() + .map_err(|_| PyValueError::new_err("User headers must be a dictionary"))?; + py_user_headers_to_rust(py, headers) + }) + .transpose()?; + let inner = RustIggyMessage::builder() + .maybe_id(id) + .payload(payload) + .maybe_user_headers(user_headers) + .build() + .map_err(to_value_error)?; Ok(Self { inner }) } } +fn to_value_error(error: impl ToString) -> PyErr { + PyValueError::new_err(error.to_string()) +} + #[derive(FromPyObject, IntoPyObject)] pub enum PyMessagePayload { #[pyo3(transparent, annotation = "str")] diff --git a/foreign/python/src/user_headers.rs b/foreign/python/src/user_headers.rs new file mode 100644 index 0000000000..68db6be65e --- /dev/null +++ b/foreign/python/src/user_headers.rs @@ -0,0 +1,857 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::BTreeMap; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; + +use iggy::prelude::{ + HeaderField, HeaderKey as RustHeaderKey, HeaderKind, HeaderValue as RustHeaderValue, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::pyclass::CompareOp; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyString}; +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods}; + +type RustUserHeaders = BTreeMap; + +#[gen_stub_pyclass_complex_enum] +#[pyclass] +/// Typed key for an Iggy user header. +/// +/// Use these constructors when the header key must preserve an explicit +/// wire type instead of using the common string-key dictionary form. +pub enum HeaderKey { + /// Raw bytes key. The byte length must be 1..=255. + Raw { value: Py }, + /// UTF-8 string key. The encoded byte length must be 1..=255. + String { value: String }, + /// Boolean key. + Bool { value: bool }, + /// Signed 8-bit integer key. + Int8 { value: i8 }, + /// Signed 16-bit integer key. + Int16 { value: i16 }, + /// Signed 32-bit integer key. + Int32 { value: i32 }, + /// Signed 64-bit integer key. + Int64 { value: i64 }, + /// Signed 128-bit integer key. + Int128 { value: i128 }, + /// Unsigned 8-bit integer key. + UnsignedInt8 { value: u8 }, + /// Unsigned 16-bit integer key. + UnsignedInt16 { value: u16 }, + /// Unsigned 32-bit integer key. + UnsignedInt32 { value: u32 }, + /// Unsigned 64-bit integer key. + UnsignedInt64 { value: u64 }, + /// Unsigned 128-bit integer key. + UnsignedInt128 { value: u128 }, + /// 32-bit floating point key. + Float32 { value: f32 }, + /// 64-bit floating point key. + Float64 { value: f64 }, +} + +#[gen_stub_pyclass_complex_enum] +#[pyclass] +/// Typed value for an Iggy user header. +/// +/// Use these constructors when the header value must preserve an explicit +/// wire type instead of using the common Python scalar dictionary form. +pub enum HeaderValue { + /// Raw bytes value. The byte length must be 1..=255. + Raw { value: Py }, + /// UTF-8 string value. The encoded byte length must be 1..=255. + String { value: String }, + /// Boolean value. + Bool { value: bool }, + /// Signed 8-bit integer value. + Int8 { value: i8 }, + /// Signed 16-bit integer value. + Int16 { value: i16 }, + /// Signed 32-bit integer value. + Int32 { value: i32 }, + /// Signed 64-bit integer value. + Int64 { value: i64 }, + /// Signed 128-bit integer value. + Int128 { value: i128 }, + /// Unsigned 8-bit integer value. + UnsignedInt8 { value: u8 }, + /// Unsigned 16-bit integer value. + UnsignedInt16 { value: u16 }, + /// Unsigned 32-bit integer value. + UnsignedInt32 { value: u32 }, + /// Unsigned 64-bit integer value. + UnsignedInt64 { value: u64 }, + /// Unsigned 128-bit integer value. + UnsignedInt128 { value: u128 }, + /// 32-bit floating point value. + Float32 { value: f32 }, + /// 64-bit floating point value. + Float64 { value: f64 }, +} + +#[gen_stub_pymethods] +#[pymethods] +impl HeaderKey { + pub fn __hash__(&self, py: Python<'_>) -> PyResult { + Ok(py_hash(header_identity_hash(self.identity(py)?))) + } + + pub fn __richcmp__( + &self, + py: Python<'_>, + other: &Bound<'_, PyAny>, + op: CompareOp, + ) -> PyResult> { + header_key_richcmp(py, self.identity(py)?, other, op) + } + + pub fn __repr__(&self, py: Python<'_>) -> PyResult { + self.repr(py) + } +} + +#[gen_stub_pymethods] +#[pymethods] +impl HeaderValue { + pub fn __hash__(&self, py: Python<'_>) -> PyResult { + Ok(py_hash(header_identity_hash(self.identity(py)?))) + } + + pub fn __richcmp__( + &self, + py: Python<'_>, + other: &Bound<'_, PyAny>, + op: CompareOp, + ) -> PyResult> { + header_value_richcmp(py, self.identity(py)?, other, op) + } + + pub fn __repr__(&self, py: Python<'_>) -> PyResult { + self.repr(py) + } +} + +impl HeaderKey { + fn to_rust(&self, py: Python<'_>) -> PyResult { + match self { + HeaderKey::Raw { value } => { + RustHeaderKey::try_from(value.extract::>(py)?).map_err(to_value_error) + } + HeaderKey::String { value } => { + RustHeaderKey::try_from(value.as_str()).map_err(to_value_error) + } + HeaderKey::Bool { value } => Ok((*value).into()), + HeaderKey::Int8 { value } => Ok((*value).into()), + HeaderKey::Int16 { value } => Ok((*value).into()), + HeaderKey::Int32 { value } => Ok((*value).into()), + HeaderKey::Int64 { value } => Ok((*value).into()), + HeaderKey::Int128 { value } => Ok((*value).into()), + HeaderKey::UnsignedInt8 { value } => Ok((*value).into()), + HeaderKey::UnsignedInt16 { value } => Ok((*value).into()), + HeaderKey::UnsignedInt32 { value } => Ok((*value).into()), + HeaderKey::UnsignedInt64 { value } => Ok((*value).into()), + HeaderKey::UnsignedInt128 { value } => Ok((*value).into()), + HeaderKey::Float32 { value } => checked_float32(*value), + HeaderKey::Float64 { value } => Ok((*value).into()), + } + } + + fn from_rust<'a>(py: Python<'a>, value: &RustHeaderKey) -> PyResult> { + let key = match value.kind() { + HeaderKind::Raw => HeaderKey::Raw { + value: PyBytes::new(py, value.as_raw().map_err(to_value_error)?).unbind(), + }, + HeaderKind::String => HeaderKey::String { + value: value.as_str().map_err(to_value_error)?.to_string(), + }, + HeaderKind::Bool => HeaderKey::Bool { + value: value.as_bool().map_err(to_value_error)?, + }, + HeaderKind::Int8 => HeaderKey::Int8 { + value: value.as_int8().map_err(to_value_error)?, + }, + HeaderKind::Int16 => HeaderKey::Int16 { + value: value.as_int16().map_err(to_value_error)?, + }, + HeaderKind::Int32 => HeaderKey::Int32 { + value: value.as_int32().map_err(to_value_error)?, + }, + HeaderKind::Int64 => HeaderKey::Int64 { + value: value.as_int64().map_err(to_value_error)?, + }, + HeaderKind::Int128 => HeaderKey::Int128 { + value: value.as_int128().map_err(to_value_error)?, + }, + HeaderKind::Uint8 => HeaderKey::UnsignedInt8 { + value: value.as_uint8().map_err(to_value_error)?, + }, + HeaderKind::Uint16 => HeaderKey::UnsignedInt16 { + value: value.as_uint16().map_err(to_value_error)?, + }, + HeaderKind::Uint32 => HeaderKey::UnsignedInt32 { + value: value.as_uint32().map_err(to_value_error)?, + }, + HeaderKind::Uint64 => HeaderKey::UnsignedInt64 { + value: value.as_uint64().map_err(to_value_error)?, + }, + HeaderKind::Uint128 => HeaderKey::UnsignedInt128 { + value: value.as_uint128().map_err(to_value_error)?, + }, + HeaderKind::Float32 => HeaderKey::Float32 { + value: value.as_float32().map_err(to_value_error)?, + }, + HeaderKind::Float64 => HeaderKey::Float64 { + value: value.as_float64().map_err(to_value_error)?, + }, + }; + key.into_pyobject(py) + } + + fn identity(&self, py: Python<'_>) -> PyResult { + match self { + HeaderKey::Raw { value } => Ok(HeaderIdentity::raw(value.extract::>(py)?)), + HeaderKey::String { value } => Ok(HeaderIdentity::string(value.as_bytes().to_vec())), + HeaderKey::Bool { value } => Ok(HeaderIdentity::bool(*value)), + HeaderKey::Int8 { value } => Ok(HeaderIdentity::int8(*value)), + HeaderKey::Int16 { value } => Ok(HeaderIdentity::int16(*value)), + HeaderKey::Int32 { value } => Ok(HeaderIdentity::int32(*value)), + HeaderKey::Int64 { value } => Ok(HeaderIdentity::int64(*value)), + HeaderKey::Int128 { value } => Ok(HeaderIdentity::int128(*value)), + HeaderKey::UnsignedInt8 { value } => Ok(HeaderIdentity::uint8(*value)), + HeaderKey::UnsignedInt16 { value } => Ok(HeaderIdentity::uint16(*value)), + HeaderKey::UnsignedInt32 { value } => Ok(HeaderIdentity::uint32(*value)), + HeaderKey::UnsignedInt64 { value } => Ok(HeaderIdentity::uint64(*value)), + HeaderKey::UnsignedInt128 { value } => Ok(HeaderIdentity::uint128(*value)), + HeaderKey::Float32 { value } => Ok(HeaderIdentity::float32(*value)), + HeaderKey::Float64 { value } => Ok(HeaderIdentity::float64(*value)), + } + } + + fn repr(&self, py: Python<'_>) -> PyResult { + match self { + HeaderKey::Raw { value } => Ok(format!( + "HeaderKey.Raw({})", + value.bind(py).repr()?.extract::()? + )), + HeaderKey::String { value } => Ok(format!("HeaderKey.String({value:?})")), + HeaderKey::Bool { value } => Ok(format!("HeaderKey.Bool({value})")), + HeaderKey::Int8 { value } => Ok(format!("HeaderKey.Int8({value})")), + HeaderKey::Int16 { value } => Ok(format!("HeaderKey.Int16({value})")), + HeaderKey::Int32 { value } => Ok(format!("HeaderKey.Int32({value})")), + HeaderKey::Int64 { value } => Ok(format!("HeaderKey.Int64({value})")), + HeaderKey::Int128 { value } => Ok(format!("HeaderKey.Int128({value})")), + HeaderKey::UnsignedInt8 { value } => Ok(format!("HeaderKey.UnsignedInt8({value})")), + HeaderKey::UnsignedInt16 { value } => Ok(format!("HeaderKey.UnsignedInt16({value})")), + HeaderKey::UnsignedInt32 { value } => Ok(format!("HeaderKey.UnsignedInt32({value})")), + HeaderKey::UnsignedInt64 { value } => Ok(format!("HeaderKey.UnsignedInt64({value})")), + HeaderKey::UnsignedInt128 { value } => Ok(format!("HeaderKey.UnsignedInt128({value})")), + HeaderKey::Float32 { value } => Ok(format!("HeaderKey.Float32({value:?})")), + HeaderKey::Float64 { value } => Ok(format!("HeaderKey.Float64({value:?})")), + } + } +} + +impl HeaderValue { + fn to_rust(&self, py: Python<'_>) -> PyResult { + match self { + HeaderValue::Raw { value } => { + RustHeaderValue::try_from(value.extract::>(py)?).map_err(to_value_error) + } + HeaderValue::String { value } => { + RustHeaderValue::try_from(value.as_str()).map_err(to_value_error) + } + HeaderValue::Bool { value } => Ok((*value).into()), + HeaderValue::Int8 { value } => Ok((*value).into()), + HeaderValue::Int16 { value } => Ok((*value).into()), + HeaderValue::Int32 { value } => Ok((*value).into()), + HeaderValue::Int64 { value } => Ok((*value).into()), + HeaderValue::Int128 { value } => Ok((*value).into()), + HeaderValue::UnsignedInt8 { value } => Ok((*value).into()), + HeaderValue::UnsignedInt16 { value } => Ok((*value).into()), + HeaderValue::UnsignedInt32 { value } => Ok((*value).into()), + HeaderValue::UnsignedInt64 { value } => Ok((*value).into()), + HeaderValue::UnsignedInt128 { value } => Ok((*value).into()), + HeaderValue::Float32 { value } => checked_float32(*value), + HeaderValue::Float64 { value } => Ok((*value).into()), + } + } + + fn from_rust<'a>(py: Python<'a>, value: &RustHeaderValue) -> PyResult> { + let value = match value.kind() { + HeaderKind::Raw => HeaderValue::Raw { + value: PyBytes::new(py, value.as_raw().map_err(to_value_error)?).unbind(), + }, + HeaderKind::String => HeaderValue::String { + value: value.as_str().map_err(to_value_error)?.to_string(), + }, + HeaderKind::Bool => HeaderValue::Bool { + value: value.as_bool().map_err(to_value_error)?, + }, + HeaderKind::Int8 => HeaderValue::Int8 { + value: value.as_int8().map_err(to_value_error)?, + }, + HeaderKind::Int16 => HeaderValue::Int16 { + value: value.as_int16().map_err(to_value_error)?, + }, + HeaderKind::Int32 => HeaderValue::Int32 { + value: value.as_int32().map_err(to_value_error)?, + }, + HeaderKind::Int64 => HeaderValue::Int64 { + value: value.as_int64().map_err(to_value_error)?, + }, + HeaderKind::Int128 => HeaderValue::Int128 { + value: value.as_int128().map_err(to_value_error)?, + }, + HeaderKind::Uint8 => HeaderValue::UnsignedInt8 { + value: value.as_uint8().map_err(to_value_error)?, + }, + HeaderKind::Uint16 => HeaderValue::UnsignedInt16 { + value: value.as_uint16().map_err(to_value_error)?, + }, + HeaderKind::Uint32 => HeaderValue::UnsignedInt32 { + value: value.as_uint32().map_err(to_value_error)?, + }, + HeaderKind::Uint64 => HeaderValue::UnsignedInt64 { + value: value.as_uint64().map_err(to_value_error)?, + }, + HeaderKind::Uint128 => HeaderValue::UnsignedInt128 { + value: value.as_uint128().map_err(to_value_error)?, + }, + HeaderKind::Float32 => HeaderValue::Float32 { + value: value.as_float32().map_err(to_value_error)?, + }, + HeaderKind::Float64 => HeaderValue::Float64 { + value: value.as_float64().map_err(to_value_error)?, + }, + }; + value.into_pyobject(py) + } + + fn identity(&self, py: Python<'_>) -> PyResult { + match self { + HeaderValue::Raw { value } => Ok(HeaderIdentity::raw(value.extract::>(py)?)), + HeaderValue::String { value } => Ok(HeaderIdentity::string(value.as_bytes().to_vec())), + HeaderValue::Bool { value } => Ok(HeaderIdentity::bool(*value)), + HeaderValue::Int8 { value } => Ok(HeaderIdentity::int8(*value)), + HeaderValue::Int16 { value } => Ok(HeaderIdentity::int16(*value)), + HeaderValue::Int32 { value } => Ok(HeaderIdentity::int32(*value)), + HeaderValue::Int64 { value } => Ok(HeaderIdentity::int64(*value)), + HeaderValue::Int128 { value } => Ok(HeaderIdentity::int128(*value)), + HeaderValue::UnsignedInt8 { value } => Ok(HeaderIdentity::uint8(*value)), + HeaderValue::UnsignedInt16 { value } => Ok(HeaderIdentity::uint16(*value)), + HeaderValue::UnsignedInt32 { value } => Ok(HeaderIdentity::uint32(*value)), + HeaderValue::UnsignedInt64 { value } => Ok(HeaderIdentity::uint64(*value)), + HeaderValue::UnsignedInt128 { value } => Ok(HeaderIdentity::uint128(*value)), + HeaderValue::Float32 { value } => Ok(HeaderIdentity::float32(*value)), + HeaderValue::Float64 { value } => Ok(HeaderIdentity::float64(*value)), + } + } + + fn repr(&self, py: Python<'_>) -> PyResult { + match self { + HeaderValue::Raw { value } => Ok(format!( + "HeaderValue.Raw({})", + value.bind(py).repr()?.extract::()? + )), + HeaderValue::String { value } => Ok(format!("HeaderValue.String({value:?})")), + HeaderValue::Bool { value } => Ok(format!("HeaderValue.Bool({value})")), + HeaderValue::Int8 { value } => Ok(format!("HeaderValue.Int8({value})")), + HeaderValue::Int16 { value } => Ok(format!("HeaderValue.Int16({value})")), + HeaderValue::Int32 { value } => Ok(format!("HeaderValue.Int32({value})")), + HeaderValue::Int64 { value } => Ok(format!("HeaderValue.Int64({value})")), + HeaderValue::Int128 { value } => Ok(format!("HeaderValue.Int128({value})")), + HeaderValue::UnsignedInt8 { value } => Ok(format!("HeaderValue.UnsignedInt8({value})")), + HeaderValue::UnsignedInt16 { value } => { + Ok(format!("HeaderValue.UnsignedInt16({value})")) + } + HeaderValue::UnsignedInt32 { value } => { + Ok(format!("HeaderValue.UnsignedInt32({value})")) + } + HeaderValue::UnsignedInt64 { value } => { + Ok(format!("HeaderValue.UnsignedInt64({value})")) + } + HeaderValue::UnsignedInt128 { value } => { + Ok(format!("HeaderValue.UnsignedInt128({value})")) + } + HeaderValue::Float32 { value } => Ok(format!("HeaderValue.Float32({value:?})")), + HeaderValue::Float64 { value } => Ok(format!("HeaderValue.Float64({value:?})")), + } + } +} + +#[derive(PartialEq, Eq, Hash)] +struct HeaderIdentity { + kind: u8, + value: Vec, +} + +impl HeaderIdentity { + fn raw(value: Vec) -> Self { + Self { + kind: HeaderKind::Raw.as_code(), + value, + } + } + + fn string(value: Vec) -> Self { + Self { + kind: HeaderKind::String.as_code(), + value, + } + } + + fn bool(value: bool) -> Self { + Self { + kind: HeaderKind::Bool.as_code(), + value: vec![u8::from(value)], + } + } + + fn int8(value: i8) -> Self { + Self { + kind: HeaderKind::Int8.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn int16(value: i16) -> Self { + Self { + kind: HeaderKind::Int16.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn int32(value: i32) -> Self { + Self { + kind: HeaderKind::Int32.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn int64(value: i64) -> Self { + Self { + kind: HeaderKind::Int64.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn int128(value: i128) -> Self { + Self { + kind: HeaderKind::Int128.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn uint8(value: u8) -> Self { + Self { + kind: HeaderKind::Uint8.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn uint16(value: u16) -> Self { + Self { + kind: HeaderKind::Uint16.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn uint32(value: u32) -> Self { + Self { + kind: HeaderKind::Uint32.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn uint64(value: u64) -> Self { + Self { + kind: HeaderKind::Uint64.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn uint128(value: u128) -> Self { + Self { + kind: HeaderKind::Uint128.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn float32(value: f32) -> Self { + Self { + kind: HeaderKind::Float32.as_code(), + value: value.to_le_bytes().to_vec(), + } + } + + fn float64(value: f64) -> Self { + Self { + kind: HeaderKind::Float64.as_code(), + value: value.to_le_bytes().to_vec(), + } + } +} + +pub(crate) fn py_user_headers_to_rust( + py: Python<'_>, + headers: &Bound<'_, PyDict>, +) -> PyResult { + // Each key/value pair is converted independently: typed `HeaderKey` / + // `HeaderValue` are used as-is, while plain Python scalars are converted on + // a best-effort basis. + let mut rust_headers = BTreeMap::new(); + for (key, value) in headers.iter() { + let key = py_header_key_to_rust(py, &key)?; + let value = py_header_value_to_rust(py, &value)?; + rust_headers.insert(key, value); + } + Ok(rust_headers) +} + +pub(crate) fn rust_user_headers_to_py<'a>( + py: Python<'a>, + headers: RustUserHeaders, +) -> PyResult> { + // Always expose the explicitly typed dict[HeaderKey, HeaderValue] so that + // no wire-type information is silently dropped. Callers who prefer the + // convenient plain form opt in through `UserHeaders.to_plain`. + let result = Bound::new(py, UserHeaders)?; + let mapping = result.as_any(); + for (key, value) in headers { + let key = HeaderKey::from_rust(py, &key)?; + let value = HeaderValue::from_rust(py, &value)?; + mapping.set_item(key, value)?; + } + Ok(result) +} + +/// User headers dictionary returned by `ReceiveMessage.user_headers`. +/// +/// This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping +/// operations work) that additionally exposes `to_plain` for the convenient +/// scalar form. +#[gen_stub_pyclass] +#[pyclass(extends=PyDict)] +pub struct UserHeaders; + +#[gen_stub_pymethods] +#[pymethods] +impl UserHeaders { + /// Wraps a mapping so its entries gain the `to_plain` helper. + /// + /// Accepts a dict whose keys and values can each independently be + /// `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool | + /// int | float`). The inherited `dict` initializer copies the provided + /// mapping. + #[new] + #[pyo3(signature = (mapping=None))] + pub fn new( + #[gen_stub(override_type(type_repr = "dict[typing.Any, typing.Any] | None"))] + mapping: Option<&Bound<'_, PyAny>>, + ) -> Self { + let _ = mapping; + UserHeaders + } + + /// Converts these headers into the convenient plain dictionary form. + /// + /// Every header kind maps losslessly onto a Python scalar, so this never + /// loses information; it only returns an error if a stored field cannot be + /// decoded. + #[gen_stub(override_return_type( + type_repr = "dict[str | bytes | bool | int | float, str | bytes | bool | int | float]" + ))] + pub fn to_plain<'a>(slf: &Bound<'a, Self>) -> PyResult> { + let py = slf.py(); + let dict = slf.as_any().cast::()?; + let headers = py_user_headers_to_rust(py, dict)?; + rust_user_headers_to_plain_py(py, headers) + } +} + +fn py_header_key_to_rust(py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult { + // A typed `HeaderKey` already carries an explicit wire type. + if let Ok(header_key) = key.extract::>() { + return header_key.to_rust(py); + } + + // Otherwise best-effort convert the plain Python scalar. + plain_to_header_field(key)?.ok_or_else(|| { + PyValueError::new_err("User header keys must be str, bytes, bool, int, float, or HeaderKey") + }) +} + +fn py_header_value_to_rust(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + // A typed `HeaderValue` already carries an explicit wire type. + if let Ok(header_value) = value.extract::>() { + return header_value.to_rust(py); + } + + // Otherwise best-effort convert the plain Python scalar. + plain_to_header_field(value)?.ok_or_else(|| { + PyValueError::new_err( + "User header values must be str, bytes, bool, int, float, or HeaderValue", + ) + }) +} + +/// Converts a plain Python scalar into a header field (used for both keys and +/// values). Returns `Ok(None)` when the object is not a supported plain type so +/// the caller can raise a message naming keys vs values; genuine range/length +/// errors are returned as `Err`. +fn plain_to_header_field(obj: &Bound<'_, PyAny>) -> PyResult>> { + // `bool` must be checked before `int` since Python bools are also ints. + if obj.is_instance_of::() { + return Ok(Some(obj.extract::()?.into())); + } + if obj.is_instance_of::() { + return int_to_header_field(obj).map(Some); + } + if obj.is_instance_of::() { + return Ok(Some(float_to_header_field(obj.extract::()?))); + } + if obj.is_instance_of::() { + return HeaderField::try_from(obj.extract::()?) + .map(Some) + .map_err(to_value_error); + } + if obj.is_instance_of::() { + return HeaderField::try_from(obj.extract::>()?) + .map(Some) + .map_err(to_value_error); + } + Ok(None) +} + +/// Picks the smallest header integer kind that represents the value exactly: +/// non-negative values use the narrowest unsigned kind, negatives the narrowest +/// signed kind. Values beyond the 128-bit range are rejected. +fn int_to_header_field(obj: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + if let Ok(value) = obj.extract::() { + return Ok(value.into()); + } + Err(PyValueError::new_err( + "User header int values must fit within the 128-bit range", + )) +} + +/// Uses `Float32` when the value is exactly representable as an f32, otherwise +/// `Float64`, so the plain float is always stored losslessly. +fn float_to_header_field(value: f64) -> HeaderField { + let narrowed = value as f32; + if f64::from(narrowed) == value { + narrowed.into() + } else { + value.into() + } +} + +fn checked_float32(value: f32) -> PyResult> { + if !value.is_finite() { + return Err(PyValueError::new_err( + "Float32 header must be a finite value within the 32-bit float range", + )); + } + Ok(value.into()) +} + +fn rust_user_headers_to_plain_py<'a>( + py: Python<'a>, + headers: RustUserHeaders, +) -> PyResult> { + // Every header kind maps losslessly onto a Python scalar, so the conversion + // is direct and only surfaces genuine decode errors. + let result = PyDict::new(py); + for (key, value) in headers { + let key = header_field_to_plain_py(py, &key)?; + let value = header_field_to_plain_py(py, &value)?; + result.set_item(key, value)?; + } + Ok(result) +} + +fn header_field_to_plain_py<'a, T>( + py: Python<'a>, + field: &HeaderField, +) -> PyResult> { + let plain = match field.kind() { + HeaderKind::Raw => PyBytes::new(py, field.as_raw().map_err(to_value_error)?).into_any(), + HeaderKind::String => field + .as_str() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Bool => field + .as_bool() + .map_err(to_value_error)? + .into_pyobject(py)? + .to_owned() + .into_any(), + HeaderKind::Int8 => field + .as_int8() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Int16 => field + .as_int16() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Int32 => field + .as_int32() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Int64 => field + .as_int64() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Int128 => field + .as_int128() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Uint8 => field + .as_uint8() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Uint16 => field + .as_uint16() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Uint32 => field + .as_uint32() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Uint64 => field + .as_uint64() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Uint128 => field + .as_uint128() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + HeaderKind::Float32 => f64::from(field.as_float32().map_err(to_value_error)?) + .into_pyobject(py)? + .into_any(), + HeaderKind::Float64 => field + .as_float64() + .map_err(to_value_error)? + .into_pyobject(py)? + .into_any(), + }; + Ok(plain) +} + +fn header_key_richcmp<'py>( + py: Python<'py>, + identity: HeaderIdentity, + other: &Bound<'py, PyAny>, + op: CompareOp, +) -> PyResult> { + let Ok(other) = other.extract::>() else { + return match op { + CompareOp::Eq => Ok(false.into_pyobject(py)?.to_owned().into_any().unbind()), + CompareOp::Ne => Ok(true.into_pyobject(py)?.to_owned().into_any().unbind()), + _ => Ok(py.NotImplemented()), + }; + }; + + let result = match op { + CompareOp::Eq => identity == other.identity(py)?, + CompareOp::Ne => identity != other.identity(py)?, + _ => return Ok(py.NotImplemented()), + }; + Ok(result.into_pyobject(py)?.to_owned().into_any().unbind()) +} + +fn header_value_richcmp<'py>( + py: Python<'py>, + identity: HeaderIdentity, + other: &Bound<'py, PyAny>, + op: CompareOp, +) -> PyResult> { + let Ok(other) = other.extract::>() else { + return match op { + CompareOp::Eq => Ok(false.into_pyobject(py)?.to_owned().into_any().unbind()), + CompareOp::Ne => Ok(true.into_pyobject(py)?.to_owned().into_any().unbind()), + _ => Ok(py.NotImplemented()), + }; + }; + + let result = match op { + CompareOp::Eq => identity == other.identity(py)?, + CompareOp::Ne => identity != other.identity(py)?, + _ => return Ok(py.NotImplemented()), + }; + Ok(result.into_pyobject(py)?.to_owned().into_any().unbind()) +} + +fn header_identity_hash(identity: HeaderIdentity) -> u64 { + let mut hasher = DefaultHasher::new(); + identity.hash(&mut hasher); + hasher.finish() +} + +fn py_hash(hash: u64) -> isize { + let hash = hash as isize; + // CPython reserves -1 as the hash error sentinel, so valid hashes with + // that value must be remapped. + if hash == -1 { -2 } else { hash } +} + +fn to_value_error(error: impl ToString) -> PyErr { + PyValueError::new_err(error.to_string()) +} diff --git a/foreign/python/tests/test_consumer_group.py b/foreign/python/tests/test_consumer_group.py index d193a9a64c..4b7a0a6f83 100644 --- a/foreign/python/tests/test_consumer_group.py +++ b/foreign/python/tests/test_consumer_group.py @@ -726,6 +726,58 @@ async def send() -> None: assert received_messages == test_messages + @pytest.mark.asyncio + async def test_consume_messages_can_read_user_headers( + self, iggy_client: IggyClient, unique_name + ): + """Test consume_messages callback receives user headers.""" + consumer_name = unique_name() + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + expected_headers: dict[str, str | bytes | bool | int | float] = { + "source": "callback", + "attempt": 1, + } + received_headers = [] + shutdown_event = asyncio.Event() + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, + name=topic_name, + partitions_count=1, + ) + + consumer = await iggy_client.consumer_group( + consumer_name, + stream_name, + topic_name, + partition_id, + PollingStrategy.Next(), + 10, + auto_commit=AutoCommit.Disabled(), + poll_interval=timedelta(milliseconds=25), + ) + + async def take(message: ReceiveMessage) -> None: + headers = message.user_headers() + assert headers is not None + received_headers.append(headers.to_plain()) + shutdown_event.set() + + async def send() -> None: + await iggy_client.send_messages( + stream_name, + topic_name, + partition_id, + [Message("callback headers", user_headers=expected_headers)], + ) + + await asyncio.gather(consumer.consume_messages(take, shutdown_event), send()) + + assert received_headers == [expected_headers] + @pytest.mark.asyncio async def test_iter_messages(self, iggy_client: IggyClient, unique_name): """Test that the consumer group can iterate over messages.""" @@ -768,6 +820,52 @@ async def test_iter_messages(self, iggy_client: IggyClient, unique_name): assert received_messages == test_messages + @pytest.mark.asyncio + async def test_iter_messages_can_read_user_headers( + self, iggy_client: IggyClient, unique_name + ): + """Test iter_messages yields messages with user headers.""" + consumer_name = unique_name() + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + expected_headers: dict[str, str | bytes | bool | int | float] = { + "source": "iterator", + "attempt": 2, + } + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, + name=topic_name, + partitions_count=1, + ) + + consumer = await iggy_client.consumer_group( + consumer_name, + stream_name, + topic_name, + partition_id, + PollingStrategy.Next(), + 10, + auto_commit=AutoCommit.Disabled(), + poll_interval=timedelta(milliseconds=25), + ) + + await iggy_client.send_messages( + stream_name, + topic_name, + partition_id, + [Message("iterator headers", user_headers=expected_headers)], + ) + + iterator = consumer.iter_messages() + message = await asyncio.wait_for(iterator.__anext__(), timeout=5) + + headers = message.user_headers() + assert headers is not None + assert headers.to_plain() == expected_headers + @pytest.mark.asyncio async def test_iter_messages_with_first_reads_existing_messages( self, iggy_client: IggyClient, unique_name diff --git a/foreign/python/tests/test_message_operations.py b/foreign/python/tests/test_message_operations.py index da57637c05..83fbcda305 100644 --- a/foreign/python/tests/test_message_operations.py +++ b/foreign/python/tests/test_message_operations.py @@ -16,13 +16,26 @@ # under the License. import asyncio +import logging import uuid +from typing import TypeAlias import pytest -from apache_iggy import IggyClient, PollingStrategy +from apache_iggy import ( + HeaderKey, + HeaderValue, + IggyClient, + PollingStrategy, + UserHeaders, +) from apache_iggy import SendMessage as Message +HTTP_CREATED = 201 +JsonValue: TypeAlias = ( + None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"] +) + class TestMessageOperations: """Test message sending, polling, and processing.""" @@ -136,8 +149,160 @@ async def test_message_properties(self, iggy_client: IggyClient, unique_name): assert isinstance(msg.offset(), int) and msg.offset() >= 0 assert isinstance(msg.id(), int) and msg.id() > 0 assert isinstance(msg.timestamp(), int) and msg.timestamp() > 0 + assert isinstance(msg.origin_timestamp(), int) and msg.origin_timestamp() > 0 assert isinstance(msg.checksum(), int) assert isinstance(msg.length(), int) and msg.length() > 0 + assert msg.user_headers() is None + + @pytest.mark.asyncio + async def test_message_user_headers_round_trip( + self, iggy_client: IggyClient, unique_name + ): + """Test plain user headers round-trip through the typed representation.""" + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + message_id = 123456789 + user_headers = { + "content-type": "application/json", + "trace-blob": b"\x00\x01", + "is-retry": False, + "attempt": 3, + "score": 0.99, + } + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, name=topic_name, partitions_count=1 + ) + + await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=partition_id, + messages=[ + Message( + "header round trip", + user_headers=user_headers, + id=message_id, + ) + ], + ) + + polled_messages = await iggy_client.poll_messages( + stream=stream_name, + topic=topic_name, + partition_id=partition_id, + polling_strategy=PollingStrategy.Last(), + count=1, + auto_commit=True, + ) + + assert len(polled_messages) == 1 + message = polled_messages[0] + assert message.id() == message_id + typed_headers = message.user_headers() + assert typed_headers is not None + assert typed_headers == { + HeaderKey.String("content-type"): HeaderValue.String("application/json"), + HeaderKey.String("trace-blob"): HeaderValue.Raw(b"\x00\x01"), + HeaderKey.String("is-retry"): HeaderValue.Bool(False), + HeaderKey.String("attempt"): HeaderValue.UnsignedInt8(3), + HeaderKey.String("score"): HeaderValue.Float64(0.99), + } + assert typed_headers.to_plain() == user_headers + assert isinstance(message.origin_timestamp(), int) + assert message.origin_timestamp() > 0 + + @pytest.mark.asyncio + async def test_typed_message_user_headers_round_trip( + self, iggy_client: IggyClient, unique_name + ): + """Test typed user headers preserve explicit header kinds.""" + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + typed_key = HeaderKey.UnsignedInt128(42) + typed_value = HeaderValue.UnsignedInt128(2**96) + user_headers: dict[HeaderKey, HeaderValue] = { + typed_key: typed_value, + HeaderKey.String("float32"): HeaderValue.Float32(1.25), + } + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, name=topic_name, partitions_count=1 + ) + + await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=partition_id, + messages=[Message("typed headers", user_headers=user_headers)], + ) + + polled_messages = await iggy_client.poll_messages( + stream=stream_name, + topic=topic_name, + partition_id=partition_id, + polling_strategy=PollingStrategy.Last(), + count=1, + auto_commit=True, + ) + + assert len(polled_messages) == 1 + headers = polled_messages[0].user_headers() + assert headers == user_headers + + @pytest.mark.asyncio + async def test_plain_scalars_pick_smallest_lossless_kind( + self, iggy_client: IggyClient, unique_name + ): + """Test plain ints/floats map to the narrowest lossless header kind.""" + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + plain_headers = { + "small": 200, + "negative": -5, + "large": 2**63, + "huge": 2**96, + "exact-float": 1.25, + "wide-float": 0.1, + } + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, name=topic_name, partitions_count=1 + ) + + await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=partition_id, + messages=[Message("plain scalars", user_headers=plain_headers)], + ) + + polled_messages = await iggy_client.poll_messages( + stream=stream_name, + topic=topic_name, + partition_id=partition_id, + polling_strategy=PollingStrategy.Last(), + count=1, + auto_commit=True, + ) + + headers = polled_messages[0].user_headers() + assert headers is not None + assert headers == { + HeaderKey.String("small"): HeaderValue.UnsignedInt8(200), + HeaderKey.String("negative"): HeaderValue.Int8(-5), + HeaderKey.String("large"): HeaderValue.UnsignedInt64(2**63), + HeaderKey.String("huge"): HeaderValue.UnsignedInt128(2**96), + HeaderKey.String("exact-float"): HeaderValue.Float32(1.25), + HeaderKey.String("wide-float"): HeaderValue.Float64(0.1), + } + assert headers.to_plain() == plain_headers @pytest.mark.asyncio @pytest.mark.parametrize( @@ -149,6 +314,115 @@ async def test_empty_payload_is_rejected(self, payload): with pytest.raises(ValueError, match="Invalid message payload length"): Message(payload) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("headers", "error"), + [ + ({"": "value"}, "Invalid header value"), + ({"x" * 256: "value"}, "Invalid header value"), + ({"key": ""}, "Invalid header value"), + ({"key": b""}, "Invalid header value"), + ({"key": "x" * 256}, "Invalid header value"), + ({object(): "value"}, "User header keys must be"), + ({"key": object()}, "User header values must be"), + ({"key": 2**128}, "128-bit range"), + ({"key": -(2**200)}, "128-bit range"), + ], + ) + async def test_invalid_user_headers_are_rejected(self, headers, error): + """Test invalid user header input raises ValueError.""" + with pytest.raises(ValueError, match=error): + Message("payload", user_headers=headers) + + def test_explicit_float32_out_of_range_is_rejected(self): + """Test an explicit Float32 whose value overflows f32 is rejected.""" + with pytest.raises(ValueError, match="32-bit float"): + Message("payload", user_headers={"k": HeaderValue.Float32(1e40)}) + + def test_mixed_typed_and_plain_headers_can_be_constructed(self): + """Test each key/value pair is converted independently and can be mixed.""" + Message( + "payload", + user_headers={ + HeaderKey.String("typed-key"): HeaderValue.UnsignedInt16(7), + HeaderKey.String("plain-value"): "still-a-string", + "plain-key": HeaderValue.Bool(True), + "fully-plain": 42, + }, + ) + + def test_typed_user_headers_can_be_constructed(self): + """Test typed header keys and values cover the full header kind surface.""" + Message( + "payload", + user_headers={ + HeaderKey.Raw(b"raw-key"): HeaderValue.Raw(b"raw-value"), + HeaderKey.String("string-key"): HeaderValue.String("string-value"), + HeaderKey.Bool(True): HeaderValue.Bool(False), + HeaderKey.Int8(-8): HeaderValue.Int8(-7), + HeaderKey.Int16(-16): HeaderValue.Int16(-15), + HeaderKey.Int32(-32): HeaderValue.Int32(-31), + HeaderKey.Int64(-64): HeaderValue.Int64(-63), + HeaderKey.Int128(-(2**80)): HeaderValue.Int128(-(2**79)), + HeaderKey.UnsignedInt8(8): HeaderValue.UnsignedInt8(9), + HeaderKey.UnsignedInt16(16): HeaderValue.UnsignedInt16(17), + HeaderKey.UnsignedInt32(32): HeaderValue.UnsignedInt32(33), + HeaderKey.UnsignedInt64(64): HeaderValue.UnsignedInt64(65), + HeaderKey.UnsignedInt128(2**80): HeaderValue.UnsignedInt128(2**79), + HeaderKey.Float32(1.25): HeaderValue.Float32(2.5), + HeaderKey.Float64(3.5): HeaderValue.Float64(4.75), + }, + ) + + def test_plain_user_headers_convert_every_kind_losslessly(self, caplog): + """Test every header kind converts to a plain scalar without logging.""" + headers = UserHeaders( + { + HeaderKey.String("content-type"): HeaderValue.String( + "application/json" + ), + HeaderKey.String("trace-blob"): HeaderValue.Raw(b"\x00\x01"), + HeaderKey.String("is-retry"): HeaderValue.Bool(True), + HeaderKey.String("attempt"): HeaderValue.Int64(3), + HeaderKey.String("schema-version"): HeaderValue.UnsignedInt16(1), + HeaderKey.String("big"): HeaderValue.UnsignedInt128(2**96), + HeaderKey.String("ratio"): HeaderValue.Float32(1.25), + HeaderKey.String("score"): HeaderValue.Float64(0.5), + } + ) + + with caplog.at_level(logging.WARNING, logger="apache_iggy"): + plain = headers.to_plain() + + assert plain == { + "content-type": "application/json", + "trace-blob": b"\x00\x01", + "is-retry": True, + "attempt": 3, + "schema-version": 1, + "big": 2**96, + "ratio": 1.25, + "score": 0.5, + } + assert caplog.records == [] + + def test_plain_user_headers_accepts_plain_dict(self): + """Test the plain dictionary form passes through unchanged.""" + headers = {"content-type": "application/json", "attempt": 3} + assert UserHeaders(headers).to_plain() == headers + + def test_plain_user_headers_preserve_non_string_keys(self, caplog): + """Test non-string typed keys convert back to their scalar Python type.""" + headers = UserHeaders( + {HeaderKey.UnsignedInt32(7): HeaderValue.String("order-id")} + ) + + with caplog.at_level(logging.WARNING, logger="apache_iggy"): + plain = headers.to_plain() + + assert plain == {7: "order-id"} + assert caplog.records == [] + @pytest.mark.asyncio @pytest.mark.parametrize( "payload", From 209ecd7e1f37df081709bc3a80421eb125506a34 Mon Sep 17 00:00:00 2001 From: Gunther Xing Date: Sun, 12 Jul 2026 18:50:58 +0800 Subject: [PATCH 2/2] chore(python): refactor Python header examples and rename to_plain to to_scalar_dict Split message-headers example into plain-headers and typed-headers variants sharing logic via common.py. Rename UserHeaders.to_plain() to to_scalar_dict() to better describe the returned dict type. --- examples/python/README.md | 19 +- .../{producer.py => common.py} | 224 +++++++++--------- examples/python/message-headers/consumer.py | 170 ------------- .../message-headers/plain-headers/consumer.py | 77 ++++++ .../message-headers/plain-headers/producer.py | 59 +++++ .../message-headers/typed-headers/consumer.py | 77 ++++++ .../message-headers/typed-headers/producer.py | 64 +++++ examples/python/pyproject.toml | 5 + foreign/python/apache_iggy.pyi | 6 +- foreign/python/src/user_headers.rs | 96 ++++++-- foreign/python/tests/test_consumer_group.py | 4 +- .../python/tests/test_message_operations.py | 27 ++- 12 files changed, 519 insertions(+), 309 deletions(-) rename examples/python/message-headers/{producer.py => common.py} (50%) delete mode 100644 examples/python/message-headers/consumer.py create mode 100644 examples/python/message-headers/plain-headers/consumer.py create mode 100644 examples/python/message-headers/plain-headers/producer.py create mode 100644 examples/python/message-headers/typed-headers/consumer.py create mode 100644 examples/python/message-headers/typed-headers/producer.py diff --git a/examples/python/README.md b/examples/python/README.md index 22fda46f25..1f3004746f 100644 --- a/examples/python/README.md +++ b/examples/python/README.md @@ -73,16 +73,25 @@ Demonstrates fundamental client connection, authentication, batch message sendin ### Message Headers -Shows how to attach and read Python SDK user headers with `str`, `bytes`, `bool`, `int`, and `float` values: +Shows how to attach and read Python SDK user headers with `str`, `bytes`, `bool`, `int`, and `float` values. Two variants share their logic through `message-headers/common.py`: + +- `plain-headers/` uses the convenient `dict[str, str | bytes | bool | int | float]` form; the SDK infers a wire type for each value. +- `typed-headers/` uses explicit `HeaderKey`/`HeaderValue` for full control over the wire type. + +Both are consumed identically, since plain headers are stored as typed headers on the wire. ```bash # Using uv -uv run message-headers/producer.py -uv run message-headers/consumer.py +uv run message-headers/plain-headers/producer.py +uv run message-headers/plain-headers/consumer.py +uv run message-headers/typed-headers/producer.py +uv run message-headers/typed-headers/consumer.py # Without using uv -python message-headers/producer.py -python message-headers/consumer.py +python message-headers/plain-headers/producer.py +python message-headers/plain-headers/consumer.py +python message-headers/typed-headers/producer.py +python message-headers/typed-headers/consumer.py ``` ## TLS Examples diff --git a/examples/python/message-headers/producer.py b/examples/python/message-headers/common.py similarity index 50% rename from examples/python/message-headers/producer.py rename to examples/python/message-headers/common.py index 7fb0b2b5e5..39e2b4de19 100644 --- a/examples/python/message-headers/producer.py +++ b/examples/python/message-headers/common.py @@ -20,10 +20,20 @@ import json import secrets import time -from collections.abc import Mapping -from typing import NamedTuple - -from apache_iggy import HeaderKey, HeaderValue, IggyClient, StreamDetails, TopicDetails +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from enum import StrEnum +from typing import Any + +from apache_iggy import ( + HeaderKey, + HeaderValue, + IggyClient, + PollingStrategy, + ReceiveMessage, + StreamDetails, + TopicDetails, +) from apache_iggy import SendMessage as Message from loguru import logger @@ -33,94 +43,28 @@ BATCHES_LIMIT = 5 MESSAGES_PER_BATCH = 10 -ORDER_CREATED_TYPE = "OrderCreated" -ORDER_CONFIRMED_TYPE = "OrderConfirmed" -ORDER_REJECTED_TYPE = "OrderRejected" - PlainHeaderValue = str | bytes | bool | int | float PlainHeaders = dict[str, PlainHeaderValue] TypedHeaders = dict[HeaderKey, HeaderValue] +HeadersBuilder = Callable[["Order"], PlainHeaders | TypedHeaders] +MessageHandler = Callable[[ReceiveMessage], None] + + +class OrderType(StrEnum): + CREATED = "OrderCreated" + CONFIRMED = "OrderConfirmed" + REJECTED = "OrderRejected" -class ArgNamespace(NamedTuple): +@dataclass(frozen=True, slots=True) +class ArgNamespace: connection_string: str -class SerializedMessage(NamedTuple): - message_type: str +@dataclass(frozen=True, slots=True) +class Order: + order_type: OrderType payload: str - headers: PlainHeaders | TypedHeaders - - -class MessagesGenerator: - def __init__(self): - self.order_id = 0 - - def generate(self) -> SerializedMessage: - self.order_id += 1 - message_type = self.order_id % 3 - - if message_type == 0: - payload = { - "orderId": f"order-{self.order_id}", - "customerId": f"customer-{secrets.randbelow(100)}", - "amount": secrets.randbelow(10000) + 1, - } - return self._serialize(ORDER_CREATED_TYPE, payload) - - if message_type == 1: - payload = { - "orderId": f"order-{self.order_id // 3}", - "timestamp": int(time.time() * 1000), - } - return self._serialize(ORDER_CONFIRMED_TYPE, payload) - - payload = { - "orderId": f"order-{self.order_id // 3}", - "reason": "Insufficient balance", - } - return self._serialize(ORDER_REJECTED_TYPE, payload) - - def _serialize( - self, message_type: str, payload: Mapping[str, object] - ) -> SerializedMessage: - encoded_payload = json.dumps(payload) - if self.order_id % 5 == 0: - # Add typed headers similar to the Rust core - typed_headers: TypedHeaders = { - HeaderKey.String("message-type"): HeaderValue.String(message_type), - HeaderKey.String("content-type"): HeaderValue.String( - "application/json" - ), - HeaderKey.String("schema-version"): HeaderValue.UnsignedInt16(1), - HeaderKey.String("created-at-ms"): HeaderValue.UnsignedInt64( - int(time.time() * 1000) - ), - HeaderKey.String("retryable"): HeaderValue.Bool( - message_type == ORDER_REJECTED_TYPE - ), - HeaderKey.String("priority-score"): HeaderValue.Float32( - (self.order_id % 100) / 100 - ), - HeaderKey.String("trace-bin"): HeaderValue.Raw( - f"trace-{self.order_id}".encode() - ), - HeaderKey.UnsignedInt32(self.order_id): HeaderValue.String("order-id"), - } - return SerializedMessage(message_type, encoded_payload, typed_headers) - - # Add easy-to-use `dict[str, str | bytes | bool | int | float]` headers - # which will be translated into typed headers by the Python SDK - plain_headers: PlainHeaders = { - "message-type": message_type, - "content-type": "application/json", - "schema-version": 1, - "created-at-ms": int(time.time() * 1000), - "retryable": message_type == ORDER_REJECTED_TYPE, - "priority-score": (self.order_id % 100) / 100, - "trace-bin": f"trace-{self.order_id}".encode(), - } - return SerializedMessage(message_type, encoded_payload, plain_headers) def parse_args() -> ArgNamespace: @@ -138,17 +82,41 @@ def parse_args() -> ArgNamespace: return ArgNamespace(**vars(parser.parse_args())) -async def main(): - args: ArgNamespace = parse_args() - client = IggyClient.from_connection_string(args.connection_string) +def generate_orders() -> Iterator[Order]: + order_id = 0 + while True: + order_id += 1 + match order_id % 3: + case 0: + payload = { + "orderId": f"order-{order_id}", + "customerId": f"customer-{secrets.randbelow(100)}", + "amount": secrets.randbelow(10000) + 1, + } + yield Order(OrderType.CREATED, json.dumps(payload)) + case 1: + payload = { + "orderId": f"order-{order_id // 3}", + "timestamp": int(time.time() * 1000), + } + yield Order(OrderType.CONFIRMED, json.dumps(payload)) + case _: + payload = { + "orderId": f"order-{order_id // 3}", + "reason": "Insufficient balance", + } + yield Order(OrderType.REJECTED, json.dumps(payload)) + + +async def connect(connection_string: str) -> IggyClient: + client = IggyClient.from_connection_string(connection_string) logger.info("Connecting to Iggy") await client.connect() logger.info("Connected") - await init_system(client) - await produce_messages(client) + return client -async def init_system(client: IggyClient): +async def init_system(client: IggyClient) -> None: try: logger.info(f"Creating stream with name {STREAM_NAME}...") stream: StreamDetails | None = await client.get_stream(STREAM_NAME) @@ -157,7 +125,6 @@ async def init_system(client: IggyClient): logger.info("Stream was created successfully.") else: logger.warning(f"Stream {stream.name} already exists with ID {stream.id}") - except Exception as error: logger.error(f"Error creating stream: {error}") logger.exception(error) @@ -180,26 +147,23 @@ async def init_system(client: IggyClient): logger.exception(error) -async def produce_messages(client: IggyClient): +async def produce_messages(client: IggyClient, build_headers: HeadersBuilder) -> None: interval = 0.5 logger.info( f"Messages will be sent to stream: {STREAM_NAME}, " f"topic: {TOPIC_NAME}, partition: {PARTITION_ID} " f"with interval {interval * 1000} ms." ) - generator = MessagesGenerator() + orders = generate_orders() sent_batches = 0 while sent_batches < BATCHES_LIMIT: messages: list[Message] = [] - for _ in range(MESSAGES_PER_BATCH): - serialized = generator.generate() - messages.append( - Message(serialized.payload, user_headers=serialized.headers) - ) + for order in (next(orders) for _ in range(MESSAGES_PER_BATCH)): + headers = build_headers(order) + messages.append(Message(order.payload, user_headers=headers)) logger.info( - f"Prepared {serialized.message_type} with headers: " - f"{format_headers(serialized.headers)}" + f"Prepared {order.order_type} with headers: {format_headers(headers)}" ) try: @@ -220,18 +184,64 @@ async def produce_messages(client: IggyClient): logger.info(f"Sent {sent_batches} batches of messages, exiting.") -def format_headers( - headers: Mapping[str, PlainHeaderValue] | Mapping[HeaderKey, HeaderValue], -) -> dict[str, str]: +async def consume_messages(client: IggyClient, handle_message: MessageHandler) -> None: + interval = 0.5 + logger.info( + f"Messages will be consumed from stream: {STREAM_NAME}, " + f"topic: {TOPIC_NAME}, partition: {PARTITION_ID} " + f"with interval {interval * 1000} ms." + ) + consumed_batches = 0 + + while consumed_batches < BATCHES_LIMIT: + try: + logger.debug("Polling for messages...") + polled_messages = await client.poll_messages( + stream=STREAM_NAME, + topic=TOPIC_NAME, + partition_id=PARTITION_ID, + polling_strategy=PollingStrategy.Next(), + count=MESSAGES_PER_BATCH, + auto_commit=True, + ) + if not polled_messages: + logger.info("No messages found in current poll") + await asyncio.sleep(interval) + continue + + for message in polled_messages: + handle_message(message) + + consumed_batches += 1 + logger.info(f"Consumed {len(polled_messages)} message(s).") + await asyncio.sleep(interval) + except Exception as error: + logger.exception(f"Exception occurred while consuming messages: {error}") + break + + logger.info(f"Consumed {consumed_batches} batches of messages, exiting.") + + +def log_order(message_type: OrderType | None, payload: object) -> None: + match message_type: + case OrderType.CREATED: + logger.info(f"Order Created: {payload}") + case OrderType.CONFIRMED: + logger.info(f"Order Confirmed: {payload}") + case OrderType.REJECTED: + logger.info(f"Order Rejected: {payload}") + case _: + logger.warning(f"Received unknown message type: {message_type}") + + +def format_headers(headers: Mapping[Any, Any]) -> dict[str, str]: formatted: dict[str, str] = {} for key, value in headers.items(): - formatted_key = repr(key) if isinstance(key, HeaderKey) else key + formatted_key = repr(key) if isinstance(key, HeaderKey) else str(key) if isinstance(value, bytes): formatted[formatted_key] = f"bytes({value.hex()})" + elif isinstance(value, HeaderValue): + formatted[formatted_key] = repr(value) else: formatted[formatted_key] = f"{value!r} ({type(value).__name__})" return formatted - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/python/message-headers/consumer.py b/examples/python/message-headers/consumer.py deleted file mode 100644 index 2f96ab84cb..0000000000 --- a/examples/python/message-headers/consumer.py +++ /dev/null @@ -1,170 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -import argparse -import asyncio -import json -from collections.abc import Mapping -from typing import Any, NamedTuple - -from apache_iggy import ( - HeaderKey, - HeaderValue, - IggyClient, - PollingStrategy, - ReceiveMessage, - UserHeaders, -) -from loguru import logger - -STREAM_NAME = "message-headers-stream" -TOPIC_NAME = "orders" -PARTITION_ID = 0 -BATCHES_LIMIT = 5 -MESSAGES_PER_BATCH = 10 - -ORDER_CREATED_TYPE = "OrderCreated" -ORDER_CONFIRMED_TYPE = "OrderConfirmed" -ORDER_REJECTED_TYPE = "OrderRejected" - - -class ArgNamespace(NamedTuple): - connection_string: str - - -def parse_args() -> ArgNamespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "connection_string", - help=( - "Connection string for Iggy client, e.g. " - "'iggy+tcp://iggy:iggy@127.0.0.1:8090'" - ), - default="iggy+tcp://iggy:iggy@127.0.0.1:8090", - nargs="?", - type=str, - ) - return ArgNamespace(**vars(parser.parse_args())) - - -async def main(): - args: ArgNamespace = parse_args() - client = IggyClient.from_connection_string(args.connection_string) - logger.info("Connecting to Iggy") - await client.connect() - logger.info("Connected") - await consume_messages(client) - - -async def consume_messages(client: IggyClient): - interval = 0.5 - logger.info( - f"Messages will be consumed from stream: {STREAM_NAME}, " - f"topic: {TOPIC_NAME}, partition: {PARTITION_ID} " - f"with interval {interval * 1000} ms." - ) - consumed_batches = 0 - - while consumed_batches < BATCHES_LIMIT: - try: - logger.debug("Polling for messages...") - polled_messages = await client.poll_messages( - stream=STREAM_NAME, - topic=TOPIC_NAME, - partition_id=PARTITION_ID, - polling_strategy=PollingStrategy.Next(), - count=MESSAGES_PER_BATCH, - auto_commit=True, - ) - if not polled_messages: - logger.info("No messages found in current poll") - await asyncio.sleep(interval) - continue - - for message in polled_messages: - handle_message(message) - - consumed_batches += 1 - logger.info(f"Consumed {len(polled_messages)} message(s).") - await asyncio.sleep(interval) - except Exception as error: - logger.exception(f"Exception occurred while consuming messages: {error}") - break - - logger.info(f"Consumed {consumed_batches} batches of messages, exiting.") - - -def handle_message(message: ReceiveMessage): - payload = json.loads(message.payload().decode("utf-8")) - # `user_headers()` returns the explicitly typed `UserHeaders` mapping - # (a dict subclass) or None when the message carries no headers. - headers = message.user_headers() - message_type = get_message_type(headers) - - logger.info( - f"Handling message at offset {message.offset()} " - f"with origin timestamp {message.origin_timestamp()}." - ) - if headers is not None: - logger.info(f"Headers: {format_headers(headers)}") - # Opt into the convenient plain form. - logger.info(f"Plain headers: {format_headers(headers.to_plain())}") - - if message_type == ORDER_CREATED_TYPE: - handle_order_created(payload) - elif message_type == ORDER_CONFIRMED_TYPE: - handle_order_confirmed(payload) - elif message_type == ORDER_REJECTED_TYPE: - handle_order_rejected(payload) - else: - logger.warning(f"Received unknown message type: {message_type}") - - -def handle_order_created(order_created: Mapping[str, Any]): - logger.info(f"Order Created: {order_created}") - - -def handle_order_confirmed(order_confirmed: Mapping[str, Any]): - logger.info(f"Order Confirmed: {order_confirmed}") - - -def handle_order_rejected(order_rejected: Mapping[str, Any]): - logger.info(f"Order Rejected: {order_rejected}") - - -def get_message_type(headers: UserHeaders | None) -> str | None: - if headers is None: - return None - message_type = headers.get(HeaderKey.String("message-type")) - if isinstance(message_type, HeaderValue.String): - return message_type.value - return None - - -def format_headers(headers: Mapping[Any, Any]) -> dict[str, str]: - formatted: dict[str, str] = {} - for key, value in headers.items(): - formatted_key = repr(key) if isinstance(key, HeaderKey) else key - if isinstance(value, bytes): - formatted[formatted_key] = f"bytes({value.hex()})" - else: - formatted[formatted_key] = f"{value!r} ({type(value).__name__})" - return formatted - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/python/message-headers/plain-headers/consumer.py b/examples/python/message-headers/plain-headers/consumer.py new file mode 100644 index 0000000000..92874b0224 --- /dev/null +++ b/examples/python/message-headers/plain-headers/consumer.py @@ -0,0 +1,77 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import asyncio +import json +import sys +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parent.parent)) + +from apache_iggy import ReceiveMessage # noqa: E402 +from common import ( # noqa: E402 + OrderType, + PlainHeaderValue, + connect, + consume_messages, + format_headers, + log_order, + parse_args, +) +from loguru import logger # noqa: E402 + + +def handle_message(message: ReceiveMessage) -> None: + payload = json.loads(message.payload().decode("utf-8")) + headers = message.user_headers() + + logger.info( + f"Handling message at offset {message.offset()} " + f"with origin timestamp {message.origin_timestamp()}." + ) + + scalar_headers: dict[PlainHeaderValue, PlainHeaderValue] = {} + if headers is not None: + # `to_scalar_dict` converts the typed headers stored on the wire back + # into the convenient plain `dict[str, str | bytes | bool | int | + # float]` form. + scalar_headers = headers.to_scalar_dict() + logger.info(f"Plain headers: {format_headers(scalar_headers)}") + + log_order(get_message_type(scalar_headers), payload) + + +def get_message_type( + headers: dict[PlainHeaderValue, PlainHeaderValue], +) -> OrderType | None: + message_type = headers.get("message-type") + if isinstance(message_type, str): + try: + return OrderType(message_type) + except ValueError: + logger.warning(f"Received unknown message type: {message_type}") + return None + + +async def main() -> None: + args = parse_args() + client = await connect(args.connection_string) + await consume_messages(client, handle_message) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/message-headers/plain-headers/producer.py b/examples/python/message-headers/plain-headers/producer.py new file mode 100644 index 0000000000..397efe2ab4 --- /dev/null +++ b/examples/python/message-headers/plain-headers/producer.py @@ -0,0 +1,59 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import asyncio +import sys +import time +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parent.parent)) + +from common import ( # noqa: E402 + Order, + OrderType, + PlainHeaders, + connect, + init_system, + parse_args, + produce_messages, +) + + +def build_plain_headers(order: Order) -> PlainHeaders: + # The plain `dict[str, str | bytes | bool | int | float]` form is the + # easiest way to attach headers: the SDK infers a wire type from each + # Python scalar and translates it into the typed form on the wire. + return { + "message-type": order.order_type, + "content-type": "application/json", + "schema-version": 1, + "created-at-ms": int(time.time() * 1000), + "retryable": order.order_type == OrderType.REJECTED, + "priority-score": 0.5, + "trace-bin": order.payload.encode(), + } + + +async def main() -> None: + args = parse_args() + client = await connect(args.connection_string) + await init_system(client) + await produce_messages(client, build_plain_headers) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/message-headers/typed-headers/consumer.py b/examples/python/message-headers/typed-headers/consumer.py new file mode 100644 index 0000000000..e24d531b21 --- /dev/null +++ b/examples/python/message-headers/typed-headers/consumer.py @@ -0,0 +1,77 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import asyncio +import json +import sys +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parent.parent)) + +from apache_iggy import ( # noqa: E402 + HeaderKey, + HeaderValue, + ReceiveMessage, + UserHeaders, +) +from common import ( # noqa: E402 + OrderType, + connect, + consume_messages, + format_headers, + log_order, + parse_args, +) +from loguru import logger # noqa: E402 + + +def handle_message(message: ReceiveMessage) -> None: + payload = json.loads(message.payload().decode("utf-8")) + # `user_headers()` returns the explicitly typed `UserHeaders` mapping + # (a dict subclass) or None when the message carries no headers. + headers = message.user_headers() + + logger.info( + f"Handling message at offset {message.offset()} " + f"with origin timestamp {message.origin_timestamp()}." + ) + if headers is not None: + logger.info(f"Headers: {format_headers(headers)}") + + log_order(get_message_type(headers), payload) + + +def get_message_type(headers: UserHeaders | None) -> OrderType | None: + if headers is None: + return None + message_type = headers.get(HeaderKey.String("message-type")) + if isinstance(message_type, HeaderValue.String): + try: + return OrderType(message_type.value) + except ValueError: + logger.warning(f"Received unknown message type: {message_type.value}") + return None + + +async def main() -> None: + args = parse_args() + client = await connect(args.connection_string) + await consume_messages(client, handle_message) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/message-headers/typed-headers/producer.py b/examples/python/message-headers/typed-headers/producer.py new file mode 100644 index 0000000000..07dac2d471 --- /dev/null +++ b/examples/python/message-headers/typed-headers/producer.py @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import asyncio +import sys +import time +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parent.parent)) + +from apache_iggy import HeaderKey, HeaderValue # noqa: E402 +from common import ( # noqa: E402 + Order, + OrderType, + TypedHeaders, + connect, + init_system, + parse_args, + produce_messages, +) + + +def build_typed_headers(order: Order) -> TypedHeaders: + # `HeaderKey`/`HeaderValue` preserve an explicit wire type. Use this form + # when you need control over the exact width/kind of a header (e.g. a + # 16-bit unsigned int) rather than letting the SDK infer it. + return { + HeaderKey.String("message-type"): HeaderValue.String(order.order_type), + HeaderKey.String("content-type"): HeaderValue.String("application/json"), + HeaderKey.String("schema-version"): HeaderValue.UnsignedInt16(1), + HeaderKey.String("created-at-ms"): HeaderValue.UnsignedInt64( + int(time.time() * 1000) + ), + HeaderKey.String("retryable"): HeaderValue.Bool( + order.order_type == OrderType.REJECTED + ), + HeaderKey.String("priority-score"): HeaderValue.Float32(0.5), + HeaderKey.String("trace-bin"): HeaderValue.Raw(order.payload.encode()), + } + + +async def main() -> None: + args = parse_args() + client = await connect(args.connection_string) + await init_system(client) + await produce_messages(client, build_typed_headers) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/pyproject.toml b/examples/python/pyproject.toml index ab6f276311..f80dfc49b7 100644 --- a/examples/python/pyproject.toml +++ b/examples/python/pyproject.toml @@ -54,3 +54,8 @@ exclude-newer = "7 days" project-includes = [ "**/*.py*", ] +search-path = [ + "message-headers", + "message-headers/plain-headers", + "message-headers/typed-headers", +] diff --git a/foreign/python/apache_iggy.pyi b/foreign/python/apache_iggy.pyi index 8e7e21ac82..186efd87de 100644 --- a/foreign/python/apache_iggy.pyi +++ b/foreign/python/apache_iggy.pyi @@ -1165,21 +1165,21 @@ class UserHeaders(dict): User headers dictionary returned by `ReceiveMessage.user_headers`. This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping - operations work) that additionally exposes `to_plain` for the convenient + operations work) that additionally exposes `to_scalar_dict` for the convenient scalar form. """ def __new__( cls, mapping: dict[typing.Any, typing.Any] | None = None ) -> UserHeaders: r""" - Wraps a mapping so its entries gain the `to_plain` helper. + Wraps a mapping so its entries gain the `to_scalar_dict` helper. Accepts a dict whose keys and values can each independently be `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool | int | float`). The inherited `dict` initializer copies the provided mapping. """ - def to_plain( + def to_scalar_dict( self, ) -> dict[str | bytes | bool | int | float, str | bytes | bool | int | float]: r""" diff --git a/foreign/python/src/user_headers.rs b/foreign/python/src/user_headers.rs index 68db6be65e..b9c1f164be 100644 --- a/foreign/python/src/user_headers.rs +++ b/foreign/python/src/user_headers.rs @@ -150,9 +150,20 @@ impl HeaderValue { } } -impl HeaderKey { - fn to_rust(&self, py: Python<'_>) -> PyResult { - match self { +/// Borrows a typed Python [`HeaderKey`] together with the GIL token so it can +/// be converted into its Rust representation (the `Raw` variant needs the token +/// to read its `PyBytes` payload). +struct PyHeaderKeyRef<'py, 'value> { + py: Python<'py>, + value: &'value HeaderKey, +} + +impl TryFrom> for RustHeaderKey { + type Error = PyErr; + + fn try_from(key: PyHeaderKeyRef<'_, '_>) -> PyResult { + let PyHeaderKeyRef { py, value } = key; + match value { HeaderKey::Raw { value } => { RustHeaderKey::try_from(value.extract::>(py)?).map_err(to_value_error) } @@ -174,8 +185,21 @@ impl HeaderKey { HeaderKey::Float64 { value } => Ok((*value).into()), } } +} - fn from_rust<'a>(py: Python<'a>, value: &RustHeaderKey) -> PyResult> { +/// Borrows a Rust [`RustHeaderKey`] together with the GIL token so it can be +/// converted into a typed Python [`HeaderKey`] (the `Raw` variant needs the +/// token to allocate its `PyBytes` payload). +struct RustHeaderKeyRef<'py, 'value> { + py: Python<'py>, + value: &'value RustHeaderKey, +} + +impl<'py> TryFrom> for Bound<'py, HeaderKey> { + type Error = PyErr; + + fn try_from(key: RustHeaderKeyRef<'py, '_>) -> PyResult { + let RustHeaderKeyRef { py, value } = key; let key = match value.kind() { HeaderKind::Raw => HeaderKey::Raw { value: PyBytes::new(py, value.as_raw().map_err(to_value_error)?).unbind(), @@ -225,7 +249,9 @@ impl HeaderKey { }; key.into_pyobject(py) } +} +impl HeaderKey { fn identity(&self, py: Python<'_>) -> PyResult { match self { HeaderKey::Raw { value } => Ok(HeaderIdentity::raw(value.extract::>(py)?)), @@ -270,9 +296,20 @@ impl HeaderKey { } } -impl HeaderValue { - fn to_rust(&self, py: Python<'_>) -> PyResult { - match self { +/// Borrows a typed Python [`HeaderValue`] together with the GIL token so it can +/// be converted into its Rust representation (the `Raw` variant needs the token +/// to read its `PyBytes` payload). +struct PyHeaderValueRef<'py, 'value> { + py: Python<'py>, + value: &'value HeaderValue, +} + +impl TryFrom> for RustHeaderValue { + type Error = PyErr; + + fn try_from(value: PyHeaderValueRef<'_, '_>) -> PyResult { + let PyHeaderValueRef { py, value } = value; + match value { HeaderValue::Raw { value } => { RustHeaderValue::try_from(value.extract::>(py)?).map_err(to_value_error) } @@ -294,8 +331,21 @@ impl HeaderValue { HeaderValue::Float64 { value } => Ok((*value).into()), } } +} - fn from_rust<'a>(py: Python<'a>, value: &RustHeaderValue) -> PyResult> { +/// Borrows a Rust [`RustHeaderValue`] together with the GIL token so it can be +/// converted into a typed Python [`HeaderValue`] (the `Raw` variant needs the +/// token to allocate its `PyBytes` payload). +struct RustHeaderValueRef<'py, 'value> { + py: Python<'py>, + value: &'value RustHeaderValue, +} + +impl<'py> TryFrom> for Bound<'py, HeaderValue> { + type Error = PyErr; + + fn try_from(value: RustHeaderValueRef<'py, '_>) -> PyResult { + let RustHeaderValueRef { py, value } = value; let value = match value.kind() { HeaderKind::Raw => HeaderValue::Raw { value: PyBytes::new(py, value.as_raw().map_err(to_value_error)?).unbind(), @@ -345,7 +395,9 @@ impl HeaderValue { }; value.into_pyobject(py) } +} +impl HeaderValue { fn identity(&self, py: Python<'_>) -> PyResult { match self { HeaderValue::Raw { value } => Ok(HeaderIdentity::raw(value.extract::>(py)?)), @@ -522,7 +574,11 @@ pub(crate) fn py_user_headers_to_rust( for (key, value) in headers.iter() { let key = py_header_key_to_rust(py, &key)?; let value = py_header_value_to_rust(py, &value)?; - rust_headers.insert(key, value); + if rust_headers.insert(key, value).is_some() { + return Err(PyValueError::new_err( + "Duplicate user header key: each header key must be unique", + )); + } } Ok(rust_headers) } @@ -533,12 +589,12 @@ pub(crate) fn rust_user_headers_to_py<'a>( ) -> PyResult> { // Always expose the explicitly typed dict[HeaderKey, HeaderValue] so that // no wire-type information is silently dropped. Callers who prefer the - // convenient plain form opt in through `UserHeaders.to_plain`. + // convenient plain form opt in through `UserHeaders.to_scalar_dict`. let result = Bound::new(py, UserHeaders)?; let mapping = result.as_any(); for (key, value) in headers { - let key = HeaderKey::from_rust(py, &key)?; - let value = HeaderValue::from_rust(py, &value)?; + let key = Bound::::try_from(RustHeaderKeyRef { py, value: &key })?; + let value = Bound::::try_from(RustHeaderValueRef { py, value: &value })?; mapping.set_item(key, value)?; } Ok(result) @@ -547,7 +603,7 @@ pub(crate) fn rust_user_headers_to_py<'a>( /// User headers dictionary returned by `ReceiveMessage.user_headers`. /// /// This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping -/// operations work) that additionally exposes `to_plain` for the convenient +/// operations work) that additionally exposes `to_scalar_dict` for the convenient /// scalar form. #[gen_stub_pyclass] #[pyclass(extends=PyDict)] @@ -556,7 +612,7 @@ pub struct UserHeaders; #[gen_stub_pymethods] #[pymethods] impl UserHeaders { - /// Wraps a mapping so its entries gain the `to_plain` helper. + /// Wraps a mapping so its entries gain the `to_scalar_dict` helper. /// /// Accepts a dict whose keys and values can each independently be /// `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool | @@ -580,7 +636,7 @@ impl UserHeaders { #[gen_stub(override_return_type( type_repr = "dict[str | bytes | bool | int | float, str | bytes | bool | int | float]" ))] - pub fn to_plain<'a>(slf: &Bound<'a, Self>) -> PyResult> { + pub fn to_scalar_dict<'a>(slf: &Bound<'a, Self>) -> PyResult> { let py = slf.py(); let dict = slf.as_any().cast::()?; let headers = py_user_headers_to_rust(py, dict)?; @@ -591,7 +647,10 @@ impl UserHeaders { fn py_header_key_to_rust(py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult { // A typed `HeaderKey` already carries an explicit wire type. if let Ok(header_key) = key.extract::>() { - return header_key.to_rust(py); + return RustHeaderKey::try_from(PyHeaderKeyRef { + py, + value: &header_key, + }); } // Otherwise best-effort convert the plain Python scalar. @@ -603,7 +662,10 @@ fn py_header_key_to_rust(py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult, value: &Bound<'_, PyAny>) -> PyResult { // A typed `HeaderValue` already carries an explicit wire type. if let Ok(header_value) = value.extract::>() { - return header_value.to_rust(py); + return RustHeaderValue::try_from(PyHeaderValueRef { + py, + value: &header_value, + }); } // Otherwise best-effort convert the plain Python scalar. diff --git a/foreign/python/tests/test_consumer_group.py b/foreign/python/tests/test_consumer_group.py index 4b7a0a6f83..836293f045 100644 --- a/foreign/python/tests/test_consumer_group.py +++ b/foreign/python/tests/test_consumer_group.py @@ -763,7 +763,7 @@ async def test_consume_messages_can_read_user_headers( async def take(message: ReceiveMessage) -> None: headers = message.user_headers() assert headers is not None - received_headers.append(headers.to_plain()) + received_headers.append(headers.to_scalar_dict()) shutdown_event.set() async def send() -> None: @@ -864,7 +864,7 @@ async def test_iter_messages_can_read_user_headers( headers = message.user_headers() assert headers is not None - assert headers.to_plain() == expected_headers + assert headers.to_scalar_dict() == expected_headers @pytest.mark.asyncio async def test_iter_messages_with_first_reads_existing_messages( diff --git a/foreign/python/tests/test_message_operations.py b/foreign/python/tests/test_message_operations.py index 83fbcda305..55368eeff2 100644 --- a/foreign/python/tests/test_message_operations.py +++ b/foreign/python/tests/test_message_operations.py @@ -210,7 +210,7 @@ async def test_message_user_headers_round_trip( HeaderKey.String("attempt"): HeaderValue.UnsignedInt8(3), HeaderKey.String("score"): HeaderValue.Float64(0.99), } - assert typed_headers.to_plain() == user_headers + assert typed_headers.to_scalar_dict() == user_headers assert isinstance(message.origin_timestamp(), int) assert message.origin_timestamp() > 0 @@ -302,7 +302,7 @@ async def test_plain_scalars_pick_smallest_lossless_kind( HeaderKey.String("exact-float"): HeaderValue.Float32(1.25), HeaderKey.String("wide-float"): HeaderValue.Float64(0.1), } - assert headers.to_plain() == plain_headers + assert headers.to_scalar_dict() == plain_headers @pytest.mark.asyncio @pytest.mark.parametrize( @@ -334,6 +334,23 @@ async def test_invalid_user_headers_are_rejected(self, headers, error): with pytest.raises(ValueError, match=error): Message("payload", user_headers=headers) + def test_duplicate_user_header_keys_are_rejected(self): + """Test a message with duplicate user header keys is rejected.""" + with pytest.raises(ValueError, match="Duplicate user header key"): + Message( + "payload", + user_headers={ + HeaderKey.String("dup"): HeaderValue.String("first"), + "dup": "second", + }, + ) + + def test_user_headers_over_100k_bytes_are_rejected(self): + """Test user headers exceeding the 100 KB limit are rejected.""" + oversized_headers = {f"key-{index:05d}": "v" * 255 for index in range(1000)} + with pytest.raises(ValueError, match="Too big headers payload"): + Message("payload", user_headers=oversized_headers) + def test_explicit_float32_out_of_range_is_rejected(self): """Test an explicit Float32 whose value overflows f32 is rejected.""" with pytest.raises(ValueError, match="32-bit float"): @@ -392,7 +409,7 @@ def test_plain_user_headers_convert_every_kind_losslessly(self, caplog): ) with caplog.at_level(logging.WARNING, logger="apache_iggy"): - plain = headers.to_plain() + plain = headers.to_scalar_dict() assert plain == { "content-type": "application/json", @@ -409,7 +426,7 @@ def test_plain_user_headers_convert_every_kind_losslessly(self, caplog): def test_plain_user_headers_accepts_plain_dict(self): """Test the plain dictionary form passes through unchanged.""" headers = {"content-type": "application/json", "attempt": 3} - assert UserHeaders(headers).to_plain() == headers + assert UserHeaders(headers).to_scalar_dict() == headers def test_plain_user_headers_preserve_non_string_keys(self, caplog): """Test non-string typed keys convert back to their scalar Python type.""" @@ -418,7 +435,7 @@ def test_plain_user_headers_preserve_non_string_keys(self, caplog): ) with caplog.at_level(logging.WARNING, logger="apache_iggy"): - plain = headers.to_plain() + plain = headers.to_scalar_dict() assert plain == {7: "order-id"} assert caplog.records == []