diff --git a/examples/python/README.md b/examples/python/README.md index 7e5da180d4..f3ebf7ac39 100644 --- a/examples/python/README.md +++ b/examples/python/README.md @@ -11,7 +11,7 @@ To run any example, first start the server with docker run --rm -p 8080:8080 -p 3000:3000 -p 8090:8090 apache/iggy:latest # Or build from source (recommended for development) -cd ../../ && cargo run --bin iggy-server +cd ../../ && cargo run --bin iggy-server -- --with-default-root-credentials ``` For server configuration options and help: diff --git a/examples/python/basic/producer.py b/examples/python/basic/producer.py index cf3e4b84a1..13a4ae34ba 100644 --- a/examples/python/basic/producer.py +++ b/examples/python/basic/producer.py @@ -19,7 +19,7 @@ import asyncio from typing import NamedTuple -from apache_iggy import IggyClient, StreamDetails, TopicDetails +from apache_iggy import IggyClient, Partitioning, StreamDetails, TopicDetails from apache_iggy import SendMessage as Message from loguru import logger @@ -106,7 +106,11 @@ async def produce_messages(client: IggyClient): await client.send_messages( stream=STREAM_NAME, topic=TOPIC_NAME, - partitioning=PARTITION_ID, + # A fixed strategy sends the whole batch to this partition. + # For topics with multiple partitions, Partitioning.balanced() + # distributes batches round-robin, while + # Partitioning.messages_key(key) keeps equal keys on the same partition. + partitioning=Partitioning.partition_id(PARTITION_ID), messages=messages, ) n_sent_batches += 1 diff --git a/foreign/python/apache_iggy.pyi b/foreign/python/apache_iggy.pyi index 254b23f41c..da501c471d 100644 --- a/foreign/python/apache_iggy.pyi +++ b/foreign/python/apache_iggy.pyi @@ -43,6 +43,7 @@ __all__ = [ "MaxTopicSize", "OptionSpec", "Partition", + "Partitioning", "Permissions", "PollingStrategy", "ReceiveMessage", @@ -872,6 +873,27 @@ class IggyClient: Sends a ping request to the server to check connectivity. Raises `RuntimeError` if the connection fails. """ + def describe_options( + self, scope: builtins.str + ) -> collections.abc.Awaitable[list[OptionSpec]]: + r""" + Describe the option catalog for a resource scope. + + This is the discovery surface for the `options` argument on + `create_topic`/`update_topic`: a key outside the catalog is refused at + create, and the binary transports carry only the error code back. + + Args: + scope: One of `"topic"`, `"stream"`, `"user"`. + + Returns: + An awaitable that resolves to `list[OptionSpec]`, empty for a scope + with no keys yet. + + Raises: + ValueError: If the scope name is not one of the three above. + RuntimeError: If the request fails. + """ def login_user( self, username: builtins.str, password: builtins.str ) -> collections.abc.Awaitable[None]: @@ -1034,27 +1056,6 @@ class IggyClient: Returns the stream details, or `None` if the stream does not exist. Raises `RuntimeError` on failure. """ - def describe_options( - self, scope: builtins.str - ) -> collections.abc.Awaitable[builtins.list[OptionSpec]]: - r""" - Describe the option catalog for a resource scope. - - This is the discovery surface for the `options` argument on - `create_topic`/`update_topic`: a key outside the catalog is refused at - create, and the binary transports carry only the error code back. - - Args: - scope: One of `"topic"`, `"stream"`, `"user"`. - - Returns: - An awaitable that resolves to `list[OptionSpec]`, empty for a scope - with no keys yet. - - Raises: - ValueError: If the scope name is not one of the three above. - RuntimeError: If the request fails. - """ def create_topic( self, stream: builtins.str | builtins.int, @@ -1331,15 +1332,32 @@ class IggyClient: self, stream: builtins.str | builtins.int, topic: builtins.str | builtins.int, - partitioning: builtins.int, + partitioning: Partitioning | builtins.int, messages: list[SendMessage], ) -> collections.abc.Awaitable[SendMessagesResponse]: r""" - Sends a list of messages to the specified topic. - Returns a SendMessagesResponse carrying the per-partition commit - confirmations, or a PyRuntimeError on failure. The confirmation list is - empty when the server reports no offsets, and the legacy server never - reports any. + Sends a batch of messages to a topic using the selected partitioning strategy. + + Args: + stream: Stream identifier as `str | int`. + topic: Topic identifier as `str | int`. + partitioning: A `Partitioning` strategy or an integer partition ID. + Use `Partitioning.balanced()`, `Partitioning.partition_id(id)`, or + `Partitioning.messages_key(key)`. An integer is shorthand for + `Partitioning.partition_id(id)`. + messages: Messages to send as `list[SendMessage]`. + + Returns: + An awaitable that resolves to `SendMessagesResponse`. Its confirmations + report the committed partition and batch base offset. The list is empty + when the server reports no offsets, including on the legacy server. + + Raises: + ValueError: If a string stream or topic identifier is invalid. + TypeError: If `partitioning` or `messages` has an unsupported type. + OverflowError: If a numeric stream, topic, or partition ID is outside + the supported unsigned 32-bit range. + RuntimeError: If the request fails. """ def poll_messages( self, @@ -1639,6 +1657,37 @@ class Partition: The number of messages in the partition. """ +@typing.final +class Partitioning: + r""" + Defines how a batch of messages is assigned to a topic partition. + """ + @staticmethod + def balanced() -> Partitioning: + r""" + Routes the batch to one partition selected by round-robin. + """ + @staticmethod + def partition_id(partition_id: builtins.int) -> Partitioning: + r""" + Routes the batch to the specified partition. + + `partition_id` must be between 0 and `2**32 - 1`. The topic must contain + that partition when the batch is sent. + """ + @staticmethod + def messages_key(key: builtins.str | bytes) -> Partitioning: + r""" + Routes the batch to one partition selected by hashing `key`. + + `key` may be `str` or `bytes`. Strings are encoded as UTF-8; the encoded + key must contain between 1 and 255 bytes. + + Raises: + ValueError: If the encoded key is empty or exceeds 255 bytes. + TypeError: If `key` is not `str` or `bytes`. + """ + @typing.final class Permissions: r""" @@ -2103,8 +2152,8 @@ class Topic: r""" Options admission resolved for the keys the client did not send. - Same shape as `options`. These would have resolved differently under - another server configuration. + Same shape as [`Self::options`]. These would have resolved differently + under another server configuration. """ @typing.final @@ -2168,8 +2217,8 @@ class TopicDetails: r""" Options admission resolved for the keys the client did not send. - Same shape as `options`. These would have resolved differently under - another server configuration. + Same shape as [`Self::options`]. These would have resolved differently + under another server configuration. """ @property def partitions(self) -> builtins.list[Partition]: diff --git a/foreign/python/src/client.rs b/foreign/python/src/client.rs index f669a87484..f975e0e6d6 100644 --- a/foreign/python/src/client.rs +++ b/foreign/python/src/client.rs @@ -19,7 +19,7 @@ use bytes::Bytes; use iggy::prelude::{ AutoCommit as RustAutoCommit, Consumer as RustConsumer, IggyClient as RustIggyClient, IggyExpiry as RustIggyExpiry, IggyMessage as RustMessage, MaxTopicSize as RustMaxTopicSize, - PollingStrategy as RustPollingStrategy, *, + Partitioning as RustPartitioning, PollingStrategy as RustPollingStrategy, *, }; use pyo3::PyRef; use pyo3::prelude::*; @@ -39,6 +39,7 @@ use crate::consumer::{ use crate::duration::{py_delta_to_iggy_duration, reject_zero}; use crate::identifier::PyIdentifier; use crate::options::OptionSpec as PyOptionSpec; +use crate::partitioning::PyPartitioning; use crate::permissions::Permissions as PyPermissions; use crate::receive_message::{PollingStrategy, ReceiveMessage}; use crate::send_message::{SendMessage, SendMessagesResponse as PySendMessagesResponse}; @@ -995,18 +996,36 @@ impl IggyClient { }) } - /// Sends a list of messages to the specified topic. - /// Returns a SendMessagesResponse carrying the per-partition commit - /// confirmations, or a PyRuntimeError on failure. The confirmation list is - /// empty when the server reports no offsets, and the legacy server never - /// reports any. + /// Sends a batch of messages to a topic using the selected partitioning strategy. + /// + /// Args: + /// stream: Stream identifier as `str | int`. + /// topic: Topic identifier as `str | int`. + /// partitioning: A `Partitioning` strategy or an integer partition ID. + /// Use `Partitioning.balanced()`, `Partitioning.partition_id(id)`, or + /// `Partitioning.messages_key(key)`. An integer is shorthand for + /// `Partitioning.partition_id(id)`. + /// messages: Messages to send as `list[SendMessage]`. + /// + /// Returns: + /// An awaitable that resolves to `SendMessagesResponse`. Its confirmations + /// report the committed partition and batch base offset. The list is empty + /// when the server reports no offsets, including on the legacy server. + /// + /// Raises: + /// ValueError: If a string stream or topic identifier is invalid. + /// TypeError: If `partitioning` or `messages` has an unsupported type. + /// OverflowError: If a numeric stream, topic, or partition ID is outside + /// the supported unsigned 32-bit range. + /// RuntimeError: If the request fails. #[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[SendMessagesResponse]", imports=("collections.abc")))] fn send_messages<'a>( &self, py: Python<'a>, stream: PyIdentifier, topic: PyIdentifier, - partitioning: u32, + #[gen_stub(override_type(type_repr = "Partitioning | builtins.int"))] + partitioning: PyPartitioning, #[gen_stub(override_type(type_repr = "list[SendMessage]"))] messages: &Bound<'_, PyList>, ) -> PyResult> { let messages: Vec = messages @@ -1023,7 +1042,7 @@ impl IggyClient { let stream = Identifier::try_from(stream)?; let topic = Identifier::try_from(topic)?; - let partitioning = Partitioning::partition_id(partitioning); + let partitioning = RustPartitioning::from(partitioning); let inner = self.inner.clone(); future_into_py(py, async move { diff --git a/foreign/python/src/lib.rs b/foreign/python/src/lib.rs index 5f2e128264..63facbce18 100644 --- a/foreign/python/src/lib.rs +++ b/foreign/python/src/lib.rs @@ -21,6 +21,7 @@ mod consumer; mod duration; mod identifier; mod options; +mod partitioning; mod permissions; mod receive_message; mod send_message; @@ -36,6 +37,7 @@ use consumer::{ ConsumerGroupMember, IggyConsumer, ReceiveMessageIterator, }; use options::OptionSpec; +use partitioning::Partitioning; use permissions::{GlobalPermissions, Permissions, StreamPermissions, TopicPermissions}; use pyo3::prelude::*; use receive_message::{PollingStrategy, ReceiveMessage}; @@ -62,6 +64,7 @@ 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::()?; m.add_class::()?; diff --git a/foreign/python/src/partitioning.rs b/foreign/python/src/partitioning.rs new file mode 100644 index 0000000000..98e3ce86af --- /dev/null +++ b/foreign/python/src/partitioning.rs @@ -0,0 +1,100 @@ +// 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 iggy::prelude::Partitioning as RustPartitioning; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyBytes}; +use pyo3_stub_gen::{ + derive::{gen_stub_pyclass, gen_stub_pymethods}, + impl_stub_type, +}; + +/// Defines how a batch of messages is assigned to a topic partition. +#[derive(Clone)] +#[pyclass(from_py_object)] +#[gen_stub_pyclass] +pub struct Partitioning { + pub(crate) inner: RustPartitioning, +} + +#[gen_stub_pymethods] +#[pymethods] +impl Partitioning { + /// Routes the batch to one partition selected by round-robin. + #[staticmethod] + pub fn balanced() -> Self { + Self { + inner: RustPartitioning::balanced(), + } + } + + /// Routes the batch to the specified partition. + /// + /// `partition_id` must be between 0 and `2**32 - 1`. The topic must contain + /// that partition when the batch is sent. + #[staticmethod] + pub fn partition_id(partition_id: u32) -> Self { + Self { + inner: RustPartitioning::partition_id(partition_id), + } + } + + /// Routes the batch to one partition selected by hashing `key`. + /// + /// `key` may be `str` or `bytes`. Strings are encoded as UTF-8; the encoded + /// key must contain between 1 and 255 bytes. + /// + /// Raises: + /// ValueError: If the encoded key is empty or exceeds 255 bytes. + /// TypeError: If `key` is not `str` or `bytes`. + #[staticmethod] + pub fn messages_key(py: Python<'_>, key: PyMessagesKey) -> PyResult { + let key = match key { + PyMessagesKey::String(key) => key.into_bytes(), + PyMessagesKey::Bytes(key) => key.extract::>(py)?, + }; + let inner = RustPartitioning::messages_key(&key) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(Self { inner }) + } +} + +#[derive(FromPyObject)] +pub enum PyMessagesKey { + #[pyo3(transparent, annotation = "str")] + String(String), + #[pyo3(transparent, annotation = "bytes")] + Bytes(Py), +} +impl_stub_type!(PyMessagesKey = String | PyBytes); + +#[derive(FromPyObject)] +pub(crate) enum PyPartitioning { + #[pyo3(transparent, annotation = "Partitioning")] + Strategy(Partitioning), + #[pyo3(transparent, annotation = "int")] + PartitionId(u32), +} +impl_stub_type!(PyPartitioning = Partitioning | isize); + +impl From for RustPartitioning { + fn from(partitioning: PyPartitioning) -> Self { + match partitioning { + PyPartitioning::Strategy(partitioning) => partitioning.inner, + PyPartitioning::PartitionId(partition_id) => Self::partition_id(partition_id), + } + } +} diff --git a/foreign/python/tests/test_message_operations.py b/foreign/python/tests/test_message_operations.py index bbf5b1d0e8..af5924c4dc 100644 --- a/foreign/python/tests/test_message_operations.py +++ b/foreign/python/tests/test_message_operations.py @@ -25,12 +25,54 @@ HeaderKey, HeaderValue, IggyClient, + Partitioning, PollingStrategy, UserHeaders, ) from apache_iggy import SendMessage as Message +class TestPartitioning: + """Test message partitioning strategy construction.""" + + @pytest.mark.unit + def test_balanced_and_partition_id_strategies(self): + assert isinstance(Partitioning.balanced(), Partitioning) + assert isinstance(Partitioning.partition_id(1), Partitioning) + assert isinstance(Partitioning.partition_id(2**32 - 1), Partitioning) + + @pytest.mark.unit + @pytest.mark.parametrize("key", [b"customer-42", "customer-42", "客户-42"]) + def test_messages_key_accepts_bytes_and_strings(self, key): + assert isinstance(Partitioning.messages_key(key), Partitioning) + + @pytest.mark.unit + @pytest.mark.parametrize("key", [b"a" * 255, "a" * 255, "界" * 85]) + def test_messages_key_accepts_255_bytes(self, key): + assert isinstance(Partitioning.messages_key(key), Partitioning) + + @pytest.mark.unit + @pytest.mark.parametrize("key", [b"", "", b"a" * 256, "a" * 256, "界" * 86]) + def test_messages_key_rejects_invalid_encoded_length(self, key): + with pytest.raises(ValueError): + Partitioning.messages_key(key) + + @pytest.mark.unit + @pytest.mark.parametrize("partition_id", [-1, 2**32]) + def test_partition_id_rejects_values_outside_u32(self, partition_id): + with pytest.raises(OverflowError): + Partitioning.partition_id(partition_id) + + @pytest.mark.unit + def test_partitioning_rejects_invalid_types(self): + with pytest.raises(TypeError): + # pyrefly: ignore # bad-argument-type + Partitioning.partition_id("0") + with pytest.raises(TypeError): + # pyrefly: ignore # bad-argument-type + Partitioning.messages_key(1) + + class TestMessageOperations: """Test message sending, polling, and processing.""" @@ -112,6 +154,101 @@ async def test_send_messages_reports_committed_confirmation( assert confirmation.partition_id == partition_id assert confirmation.base_offset == polled_messages[0].offset() + @pytest.mark.asyncio + async def test_send_messages_with_partition_id_strategy( + self, iggy_client: IggyClient, unique_name + ): + stream_name = unique_name() + topic_name = unique_name() + partition_id = 2 + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, name=topic_name, partitions_count=3 + ) + + response = await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=Partitioning.partition_id(partition_id), + messages=[Message("fixed partition 1"), Message("fixed partition 2")], + ) + + assert len(response.confirmations) == 1 + assert response.confirmations[0].partition_id == partition_id + + with pytest.raises(RuntimeError): + await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=Partitioning.partition_id(3), + messages=[Message("missing partition")], + ) + + @pytest.mark.asyncio + async def test_send_messages_with_balanced_strategy( + self, iggy_client: IggyClient, unique_name + ): + stream_name = unique_name() + topic_name = unique_name() + partitions_count = 3 + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, + name=topic_name, + partitions_count=partitions_count, + ) + + responses = [ + await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=Partitioning.balanced(), + messages=[Message(f"balanced {index}")], + ) + for index in range(partitions_count) + ] + + assert all(len(response.confirmations) == 1 for response in responses) + assert ( + len({response.confirmations[0].partition_id for response in responses}) + == partitions_count + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("key", [b"customer-42", "customer-42"]) + async def test_send_messages_with_same_key_uses_same_partition( + self, iggy_client: IggyClient, unique_name, key + ): + stream_name = unique_name() + topic_name = unique_name() + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, name=topic_name, partitions_count=3 + ) + partitioning = Partitioning.messages_key(key) + + first = await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=partitioning, + messages=[Message("first")], + ) + second = await iggy_client.send_messages( + stream=stream_name, + topic=topic_name, + partitioning=partitioning, + messages=[Message("second")], + ) + + assert len(first.confirmations) == 1 + assert len(second.confirmations) == 1 + assert ( + first.confirmations[0].partition_id == second.confirmations[0].partition_id + ) + @pytest.mark.asyncio async def test_send_and_poll_messages_as_bytes( self, iggy_client: IggyClient, unique_name