ethanlin01x commented on code in PR #3927:
URL: https://github.com/apache/iggy/pull/3927#discussion_r3836019296
##########
foreign/python/apache_iggy.pyi:
##########
@@ -1601,6 +1605,30 @@ class Partition:
The number of messages in the partition.
"""
[email protected]
+class Partitioning:
Review Comment:
`Partitioning` is missing from `__all__`, so from apache_iggy import *
fails pyrefly with unknown-name.
##########
foreign/python/src/partitioning.rs:
##########
@@ -0,0 +1,93 @@
+// 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 partitions using server-side round-robin selection.
+ #[staticmethod]
+ pub fn balanced() -> Self {
+ Self {
+ inner: RustPartitioning::balanced(),
+ }
+ }
+
+ /// Routes the batch to the specified partition.
+ #[staticmethod]
+ pub fn partition_id(partition_id: u32) -> Self {
+ Self {
+ inner: RustPartitioning::partition_id(partition_id),
+ }
+ }
+
+ /// Routes the batch using a binary key hashed by the server.
+ ///
+ /// String keys are encoded as UTF-8. The encoded key must contain between
+ /// 1 and 255 bytes.
+ #[staticmethod]
+ pub fn messages_key(py: Python<'_>, key: PyMessagesKey) -> PyResult<Self> {
+ let key = match key {
+ PyMessagesKey::String(key) => key.into_bytes(),
+ PyMessagesKey::Bytes(key) => key.extract::<Vec<u8>>(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<PyBytes>),
+}
+impl_stub_type!(PyMessagesKey = String | PyBytes);
+
+#[derive(FromPyObject)]
+pub(crate) enum PyPartitioning {
+ #[pyo3(transparent)]
+ Strategy(Partitioning),
Review Comment:
Strategy has no annotation, so errors read 'Strategy | int' — a Rust name
that does not exist in Python. Please add annotation = "Partitioning"
##########
foreign/python/tests/test_message_operations.py:
##########
@@ -109,6 +150,87 @@ 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")],
+ )
+
+ assert len(response.confirmations) == 1
+ assert response.confirmations[0].partition_id == partition_id
+
+ @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,
+ )
+
+ response = await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=Partitioning.balanced(),
+ messages=[Message("balanced")],
+ )
+
+ assert len(response.confirmations) == 1
+ assert response.confirmations[0].partition_id < partitions_count
Review Comment:
Sending once cannot tell round-robin apart from always returning partition
0. Could you send three batches and assert the three confirmed partition ids
are distinct?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
