foreign/python/src/consumer.rs - iggy - Git at Google

 // 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::sync::Arc;
 use std::time::Duration;

 use futures::StreamExt;
 use iggy::consumer_ext::{IggyConsumerMessageExt, MessageConsumer};
 use iggy::prelude::{
     AutoCommit as RustAutoCommit, AutoCommitAfter as RustAutoCommitAfter,
     AutoCommitWhen as RustAutoCommitWhen, ConsumerGroup as RustConsumerGroup,
     ConsumerGroupDetails as RustConsumerGroupDetails,
     ConsumerGroupMember as RustConsumerGroupMember, IggyConsumer as RustIggyConsumer, IggyDuration,
     IggyError, ReceivedMessage,
 };
 use pyo3::exceptions::PyStopAsyncIteration;
 use pyo3::types::{PyDelta, PyDeltaAccess};

 use pyo3::prelude::*;
 use pyo3_async_runtimes::TaskLocals;
 use pyo3_async_runtimes::tokio::{future_into_py, get_runtime, into_future, scope};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
 use pyo3_stub_gen::{PyStubType, TypeInfo};
 use tokio::sync::Mutex;
 use tokio::sync::oneshot::Sender;
 use tokio::task::JoinHandle;

 use crate::identifier::PyIdentifier;
 use crate::receive_message::ReceiveMessage;

 /// A Python class representing the Iggy consumer.
 /// It wraps the RustIggyConsumer and provides asynchronous functionality
 /// through the contained runtime.
 #[gen_stub_pyclass]
 #[pyclass]
 pub struct IggyConsumer {
     pub(crate) inner: Arc<Mutex<RustIggyConsumer>>,
 }

 #[gen_stub_pymethods]
 #[pymethods]
 impl IggyConsumer {
     /// Get the last consumed offset or `None` if no offset has been consumed yet.
     #[gen_stub(override_return_type(type_repr = "builtins.int | None"))]
     fn get_last_consumed_offset(&self, partition_id: u32) -> Option<u64> {
         self.inner
             .blocking_lock()
             .get_last_consumed_offset(partition_id)
     }

     /// Get the last stored offset or `None` if no offset has been stored yet.
     #[gen_stub(override_return_type(type_repr = "builtins.int | None"))]
     fn get_last_stored_offset(&self, partition_id: u32) -> Option<u64> {
         self.inner
             .blocking_lock()
             .get_last_stored_offset(partition_id)
     }

     /// Gets the name of the consumer group.
     fn name(&self) -> String {
         self.inner.blocking_lock().name().to_string()
     }

     /// Gets the current partition id or `0` if no messages have been polled yet.
     fn partition_id(&self) -> u32 {
         self.inner.blocking_lock().partition_id()
     }

     /// Gets the name of the stream this consumer group is configured for.
     fn stream(&self) -> PyResult<PyIdentifier> {
         let guard = self.inner.blocking_lock();
         PyIdentifier::try_from(guard.stream())
     }

     /// Gets the name of the topic this consumer group is configured for.
     fn topic(&self) -> PyResult<PyIdentifier> {
         let guard = self.inner.blocking_lock();
         PyIdentifier::try_from(guard.topic())
     }

     /// Stores the provided offset for the provided partition id or if none is specified
     /// uses the current partition id for the consumer group.
     /// Returns `Ok(())` if the server responds successfully, or a `PyRuntimeError`
     /// if the operation fails.
     #[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
     fn store_offset<'a>(
         &self,
         py: Python<'a>,
         offset: u64,
         #[gen_stub(override_type(type_repr = "builtins.int | None"))] partition_id: Option<u32>,
     ) -> PyResult<Bound<'a, PyAny>> {
         let inner = self.inner.clone();
         future_into_py(py, async move {
             inner
                 .lock()
                 .await
                 .store_offset(offset, partition_id)
                 .await
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
         })
     }

     /// Deletes the offset for the provided partition id or if none is specified
     /// uses the current partition id for the consumer group.
     /// Returns `Ok(())` if the server responds successfully, or a `PyRuntimeError`
     /// if the operation fails.
     #[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
     fn delete_offset<'a>(
         &self,
         py: Python<'a>,
         #[gen_stub(override_type(type_repr = "builtins.int | None"))] partition_id: Option<u32>,
     ) -> PyResult<Bound<'a, PyAny>> {
         let inner = self.inner.clone();
         future_into_py(py, async move {
             inner
                 .lock()
                 .await
                 .delete_offset(partition_id)
                 .await
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
         })
     }

     /// Asynchronously iterate over `ReceiveMessage`s.
     /// Returns an async iterator that raises `StopAsyncIteration` when no more messages are available
     /// or a `PyRuntimeError` on failure.
     /// Note: This method does not currently support `AutoCommit.After`.
     /// For `AutoCommit.IntervalOrAfter(datetime.timedelta, AutoCommitAfter)`,
     /// only the interval part is applied; the `after` mode is ignored.
     /// Use `consume_messages()` if you need commit-after-processing semantics.
     #[gen_stub(override_return_type(type_repr="collections.abc.AsyncIterator[ReceiveMessage]", imports=("collections.abc")))]
     fn iter_messages(&self) -> ReceiveMessageIterator {
         let inner = self.inner.clone();
         ReceiveMessageIterator { inner }
     }

     /// Consumes messages continuously using a callback function and an optional `asyncio.Event` for signaling shutdown.
     /// Returns an awaitable that completes when shutdown is signaled or a PyRuntimeError on failure.
     #[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
     fn consume_messages<'a>(
         &self,
         py: Python<'a>,
         #[gen_stub(override_type(type_repr="collections.abc.Callable[[ReceiveMessage], collections.abc.Awaitable[None]]", imports=("collections.abc")))]
         callback: Bound<'a, PyAny>,
         #[gen_stub(override_type(type_repr="asyncio.Event | None", imports=("asyncio")))]
         shutdown_event: Option<Bound<'a, PyAny>>,
     ) -> PyResult<Bound<'a, PyAny>> {
         let inner = self.inner.clone();
         let callback: Py<PyAny> = callback.unbind();
         let shutdown_event: Option<Py<PyAny>> = shutdown_event.map(|e| e.unbind());

         future_into_py(py, async {
             let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();

             let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
             let handle_consume: JoinHandle<PyResult<Result<(), IggyError>>> =
                 get_runtime().spawn(scope(task_locals, async move {
                     let task_locals =
                         Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
                     let consumer = PyCallbackConsumer {
                         callback: Arc::new(callback),
                         task_locals: Arc::new(Mutex::new(task_locals)),
                     };
                     let mut inner = inner.lock().await;
                     Ok(inner.consume_messages(&consumer, shutdown_rx).await)
                 }));
             let consume_result;

             if let Some(shutdown_event) = shutdown_event {
                 let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
                 async fn shutdown_impl(
                     shutdown_event: Py<PyAny>,
                     shutdown_tx: Sender<()>,
                 ) -> PyResult<()> {
                     Python::attach(|py| {
                         into_future(shutdown_event.bind(py).as_any().call_method0("wait")?)
                     })?
                     .await?;
                     shutdown_tx.send(()).map_err(|_| {
                         PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                             "Failed to signal shutdown",
                         )
                     })?;
                     Ok(())
                 }
                 let handle_shutdown: JoinHandle<Result<(), PyErr>> = get_runtime().spawn(scope(
                     task_locals,
                     shutdown_impl(shutdown_event, shutdown_tx),
                 ));
                 let shutdown_result;
                 (consume_result, shutdown_result) = tokio::join!(handle_consume, handle_shutdown);
                 shutdown_result.map_err(|e| {
                     PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
                 })??;
             } else {
                 consume_result = handle_consume.await;
             }

             consume_result
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))??
                 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
             Ok(())
         })
     }
 }

 #[gen_stub_pyclass]
 #[pyclass]
 pub struct ConsumerGroup {
     pub(crate) inner: RustConsumerGroup,
 }

 impl From<RustConsumerGroup> for ConsumerGroup {
     fn from(group: RustConsumerGroup) -> Self {
         Self { inner: group }
     }
 }

 #[gen_stub_pymethods]
 #[pymethods]
 impl ConsumerGroup {
     /// Gets the unique identifier (numeric) of the consumer group.
     #[getter]
     pub fn id(&self) -> u32 {
         self.inner.id
     }

     /// Gets the name of the consumer group.
     #[getter]
     pub fn name(&self) -> String {
         self.inner.name.to_string()
     }

     /// Gets the number of partitions the consumer group is consuming.
     #[getter]
     pub fn partitions_count(&self) -> u32 {
         self.inner.partitions_count
     }

     /// Gets the number of members in the consumer group.
     #[getter]
     pub fn members_count(&self) -> u32 {
         self.inner.members_count
     }
 }

 #[gen_stub_pyclass]
 #[pyclass]
 pub struct ConsumerGroupDetails {
     pub(crate) inner: RustConsumerGroupDetails,
 }

 impl From<RustConsumerGroupDetails> for ConsumerGroupDetails {
     fn from(group: RustConsumerGroupDetails) -> Self {
         Self { inner: group }
     }
 }

 #[gen_stub_pymethods]
 #[pymethods]
 impl ConsumerGroupDetails {
     /// Gets the unique identifier (numeric) of the consumer group.
     #[getter]
     pub fn id(&self) -> u32 {
         self.inner.id
     }

     /// Gets the name of the consumer group.
     #[getter]
     pub fn name(&self) -> String {
         self.inner.name.to_string()
     }

     /// Gets the number of partitions the consumer group is consuming.
     #[getter]
     pub fn partitions_count(&self) -> u32 {
         self.inner.partitions_count
     }

     /// Gets the number of members in the consumer group.
     #[getter]
     pub fn members_count(&self) -> u32 {
         self.inner.members_count
     }

     /// Gets the collection of members in the consumer group.
     #[getter]
     pub fn members(&self) -> Vec<ConsumerGroupMember> {
         self.inner
             .members
             .iter()
             .map(ConsumerGroupMember::from)
             .collect()
     }
 }

 #[gen_stub_pyclass]
 #[pyclass]
 pub struct ConsumerGroupMember {
     pub(crate) inner: RustConsumerGroupMember,
 }

 impl From<&RustConsumerGroupMember> for ConsumerGroupMember {
     fn from(member: &RustConsumerGroupMember) -> Self {
         Self {
             inner: RustConsumerGroupMember {
                 id: member.id,
                 partitions_count: member.partitions_count,
                 partitions: member.partitions.clone(),
             },
         }
     }
 }

 #[gen_stub_pymethods]
 #[pymethods]
 impl ConsumerGroupMember {
     /// Gets the unique identifier (numeric) of the consumer group member.
     #[getter]
     pub fn id(&self) -> u32 {
         self.inner.id
     }

     /// Gets the number of partitions the consumer group member is consuming.
     #[getter]
     pub fn partitions_count(&self) -> u32 {
         self.inner.partitions_count
     }

     /// Gets the collection of partitions the consumer group member is consuming.
     #[getter]
     pub fn partitions(&self) -> Vec<u32> {
         self.inner.partitions.clone()
     }
 }

 #[pyclass]
 pub struct ReceiveMessageIterator {
     pub(crate) inner: Arc<Mutex<RustIggyConsumer>>,
 }

 #[pymethods]
 impl ReceiveMessageIterator {
     pub fn __anext__<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
         let inner = self.inner.clone();
         future_into_py(py, async move {
             let mut inner = inner.lock().await;
             if let Some(message) = inner.next().await {
                 Ok(message
                     .map(|m| ReceiveMessage {
                         inner: m.message,
                         partition_id: m.partition_id,
                     })
                     .map_err(|e| {
                         PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
                     })?)
             } else {
                 Err(PyStopAsyncIteration::new_err("No more messages"))
             }
         })
     }

     pub fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
         slf
     }
 }

 struct PyCallbackConsumer {
     callback: Arc<Py<PyAny>>,
     task_locals: Arc<Mutex<TaskLocals>>,
 }

 impl MessageConsumer for PyCallbackConsumer {
     async fn consume(&self, received: ReceivedMessage) -> Result<(), IggyError> {
         let callback = self.callback.clone();
         let task_locals = self.task_locals.clone().lock_owned().await;
         let task_locals = task_locals.clone();
         let message = ReceiveMessage {
             inner: received.message,
             partition_id: received.partition_id,
         };
         get_runtime()
             .spawn(scope(task_locals, async move {
                 Python::attach(|py| {
                     let callback = callback.bind(py);
                     let result = callback.as_any().call1((message,))?;
                     into_future(result)
                 })
             }))
             .await
             .map_err(|_| IggyError::CannotReadMessage)?
             .map_err(|_| IggyError::CannotReadMessage)?
             .await
             .map_err(|_| IggyError::CannotReadMessage)?;
         Ok(())
     }
 }

 /// The auto-commit configuration for storing the offset on the server.
 // #[derive(Debug, PartialEq, Copy, Clone)]
 #[gen_stub_pyclass_complex_enum]
 #[pyclass]
 pub enum AutoCommit {
     /// The auto-commit is disabled and the offset must be stored manually by the consumer.
     Disabled(),
     /// The auto-commit is enabled and the offset is stored on the server after a certain interval.
     Interval(Py<PyDelta>),
     /// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode when consuming the messages.
     IntervalOrWhen(Py<PyDelta>, AutoCommitWhen),
     /// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode after consuming the messages.
     IntervalOrAfter(Py<PyDelta>, AutoCommitAfter),
     /// The auto-commit is enabled and the offset is stored on the server depending on the mode when consuming the messages.
     When(AutoCommitWhen),
     /// The auto-commit is enabled and the offset is stored on the server depending on the mode after consuming the messages.
     After(AutoCommitAfter),
 }

 impl From<&AutoCommit> for RustAutoCommit {
     fn from(val: &AutoCommit) -> RustAutoCommit {
         match val {
             AutoCommit::Disabled() => RustAutoCommit::Disabled,
             AutoCommit::Interval(delta) => {
                 let duration = py_delta_to_iggy_duration(delta);
                 RustAutoCommit::Interval(duration)
             }
             AutoCommit::IntervalOrWhen(delta, when) => {
                 let duration = py_delta_to_iggy_duration(delta);
                 RustAutoCommit::IntervalOrWhen(duration, when.into())
             }
             AutoCommit::IntervalOrAfter(delta, after) => {
                 let duration = py_delta_to_iggy_duration(delta);
                 RustAutoCommit::IntervalOrAfter(duration, after.into())
             }
             AutoCommit::When(when) => RustAutoCommit::When(when.into()),
             AutoCommit::After(after) => RustAutoCommit::After(after.into()),
         }
     }
 }

 /// The auto-commit mode for storing the offset on the server.
 #[derive(Debug, PartialEq, Copy, Clone)]
 #[gen_stub_pyclass_complex_enum(skip_stub_type)]
 #[pyclass(from_py_object)]
 pub enum AutoCommitWhen {
     /// The offset is stored on the server when the messages are received.
     PollingMessages(),
     /// The offset is stored on the server when all the messages are consumed.
     ConsumingAllMessages(),
     /// The offset is stored on the server when consuming each message.
     ConsumingEachMessage(),
     /// The offset is stored on the server when consuming every Nth message.
     ConsumingEveryNthMessage(u32),
 }

 impl From<&AutoCommitWhen> for RustAutoCommitWhen {
     fn from(val: &AutoCommitWhen) -> RustAutoCommitWhen {
         match val {
             AutoCommitWhen::PollingMessages() => RustAutoCommitWhen::PollingMessages,
             AutoCommitWhen::ConsumingAllMessages() => RustAutoCommitWhen::ConsumingAllMessages,
             AutoCommitWhen::ConsumingEachMessage() => RustAutoCommitWhen::ConsumingEachMessage,
             AutoCommitWhen::ConsumingEveryNthMessage(n) => {
                 RustAutoCommitWhen::ConsumingEveryNthMessage(n.to_owned())
             }
         }
     }
 }

 impl PyStubType for AutoCommitWhen {
     fn type_output() -> TypeInfo {
         TypeInfo::unqualified("AutoCommitWhen")
     }
 }

 /// The auto-commit mode for storing the offset on the server **after** receiving the messages.
 #[derive(Debug, PartialEq, Copy, Clone)]
 #[gen_stub_pyclass_complex_enum(skip_stub_type)]
 #[pyclass(from_py_object)]
 #[allow(clippy::enum_variant_names)]
 pub enum AutoCommitAfter {
     /// The offset is stored on the server after all the messages are consumed.
     ConsumingAllMessages(),
     /// The offset is stored on the server after consuming each message.
     ConsumingEachMessage(),
     /// The offset is stored on the server after consuming every Nth message.
     ConsumingEveryNthMessage(u32),
 }

 impl From<&AutoCommitAfter> for RustAutoCommitAfter {
     fn from(val: &AutoCommitAfter) -> RustAutoCommitAfter {
         match val {
             AutoCommitAfter::ConsumingAllMessages() => RustAutoCommitAfter::ConsumingAllMessages,
             AutoCommitAfter::ConsumingEachMessage() => RustAutoCommitAfter::ConsumingEachMessage,
             AutoCommitAfter::ConsumingEveryNthMessage(n) => {
                 RustAutoCommitAfter::ConsumingEveryNthMessage(n.to_owned())
             }
         }
     }
 }

 impl PyStubType for AutoCommitAfter {
     fn type_output() -> TypeInfo {
         TypeInfo::unqualified("AutoCommitAfter")
     }
 }

 pub fn py_delta_to_iggy_duration(delta1: &Py<PyDelta>) -> IggyDuration {
     Python::attach(|py| {
         let delta = delta1.bind(py);
         let seconds = (delta.get_days() * 60 * 60 * 24 + delta.get_seconds()) as u64;
         let nanos = (delta.get_microseconds() * 1_000) as u32;
         IggyDuration::new(Duration::new(seconds, nanos))
     })
 }
	// 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::sync::Arc;
	use std::time::Duration;

	use futures::StreamExt;
	use iggy::consumer_ext::{IggyConsumerMessageExt, MessageConsumer};
	use iggy::prelude::{
	AutoCommit as RustAutoCommit, AutoCommitAfter as RustAutoCommitAfter,
	AutoCommitWhen as RustAutoCommitWhen, ConsumerGroup as RustConsumerGroup,
	ConsumerGroupDetails as RustConsumerGroupDetails,
	ConsumerGroupMember as RustConsumerGroupMember, IggyConsumer as RustIggyConsumer, IggyDuration,
	IggyError, ReceivedMessage,
	};
	use pyo3::exceptions::PyStopAsyncIteration;
	use pyo3::types::{PyDelta, PyDeltaAccess};

	use pyo3::prelude::*;
	use pyo3_async_runtimes::TaskLocals;
	use pyo3_async_runtimes::tokio::{future_into_py, get_runtime, into_future, scope};
	use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
	use pyo3_stub_gen::{PyStubType, TypeInfo};
	use tokio::sync::Mutex;
	use tokio::sync::oneshot::Sender;
	use tokio::task::JoinHandle;

	use crate::identifier::PyIdentifier;
	use crate::receive_message::ReceiveMessage;

	/// A Python class representing the Iggy consumer.
	/// It wraps the RustIggyConsumer and provides asynchronous functionality
	/// through the contained runtime.
	#[gen_stub_pyclass]
	#[pyclass]
	pub struct IggyConsumer {
	pub(crate) inner: Arc<Mutex<RustIggyConsumer>>,
	}

	#[gen_stub_pymethods]
	#[pymethods]
	impl IggyConsumer {
	/// Get the last consumed offset or `None` if no offset has been consumed yet.
	#[gen_stub(override_return_type(type_repr = "builtins.int \| None"))]
	fn get_last_consumed_offset(&self, partition_id: u32) -> Option<u64> {
	self.inner
	.blocking_lock()
	.get_last_consumed_offset(partition_id)
	}

	/// Get the last stored offset or `None` if no offset has been stored yet.
	#[gen_stub(override_return_type(type_repr = "builtins.int \| None"))]
	fn get_last_stored_offset(&self, partition_id: u32) -> Option<u64> {
	self.inner
	.blocking_lock()
	.get_last_stored_offset(partition_id)
	}

	/// Gets the name of the consumer group.
	fn name(&self) -> String {
	self.inner.blocking_lock().name().to_string()
	}

	/// Gets the current partition id or `0` if no messages have been polled yet.
	fn partition_id(&self) -> u32 {
	self.inner.blocking_lock().partition_id()
	}

	/// Gets the name of the stream this consumer group is configured for.
	fn stream(&self) -> PyResult<PyIdentifier> {
	let guard = self.inner.blocking_lock();
	PyIdentifier::try_from(guard.stream())
	}

	/// Gets the name of the topic this consumer group is configured for.
	fn topic(&self) -> PyResult<PyIdentifier> {
	let guard = self.inner.blocking_lock();
	PyIdentifier::try_from(guard.topic())
	}

	/// Stores the provided offset for the provided partition id or if none is specified
	/// uses the current partition id for the consumer group.
	/// Returns `Ok(())` if the server responds successfully, or a `PyRuntimeError`
	/// if the operation fails.
	#[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
	fn store_offset<'a>(
	&self,
	py: Python<'a>,
	offset: u64,
	#[gen_stub(override_type(type_repr = "builtins.int \| None"))] partition_id: Option<u32>,
	) -> PyResult<Bound<'a, PyAny>> {
	let inner = self.inner.clone();
	future_into_py(py, async move {
	inner
	.lock()
	.await
	.store_offset(offset, partition_id)
	.await
	.map_err(\|e\| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
	})
	}

	/// Deletes the offset for the provided partition id or if none is specified
	/// uses the current partition id for the consumer group.
	/// Returns `Ok(())` if the server responds successfully, or a `PyRuntimeError`
	/// if the operation fails.
	#[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
	fn delete_offset<'a>(
	&self,
	py: Python<'a>,
	#[gen_stub(override_type(type_repr = "builtins.int \| None"))] partition_id: Option<u32>,
	) -> PyResult<Bound<'a, PyAny>> {
	let inner = self.inner.clone();
	future_into_py(py, async move {
	inner
	.lock()
	.await
	.delete_offset(partition_id)
	.await
	.map_err(\|e\| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
	})
	}

	/// Asynchronously iterate over `ReceiveMessage`s.
	/// Returns an async iterator that raises `StopAsyncIteration` when no more messages are available
	/// or a `PyRuntimeError` on failure.
	/// Note: This method does not currently support `AutoCommit.After`.
	/// For `AutoCommit.IntervalOrAfter(datetime.timedelta, AutoCommitAfter)`,
	/// only the interval part is applied; the `after` mode is ignored.
	/// Use `consume_messages()` if you need commit-after-processing semantics.
	#[gen_stub(override_return_type(type_repr="collections.abc.AsyncIterator[ReceiveMessage]", imports=("collections.abc")))]
	fn iter_messages(&self) -> ReceiveMessageIterator {
	let inner = self.inner.clone();
	ReceiveMessageIterator { inner }
	}

	/// Consumes messages continuously using a callback function and an optional `asyncio.Event` for signaling shutdown.
	/// Returns an awaitable that completes when shutdown is signaled or a PyRuntimeError on failure.
	#[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[None]", imports=("collections.abc")))]
	fn consume_messages<'a>(
	&self,
	py: Python<'a>,
	#[gen_stub(override_type(type_repr="collections.abc.Callable[[ReceiveMessage], collections.abc.Awaitable[None]]", imports=("collections.abc")))]
	callback: Bound<'a, PyAny>,
	#[gen_stub(override_type(type_repr="asyncio.Event \| None", imports=("asyncio")))]
	shutdown_event: Option<Bound<'a, PyAny>>,
	) -> PyResult<Bound<'a, PyAny>> {
	let inner = self.inner.clone();
	let callback: Py<PyAny> = callback.unbind();
	let shutdown_event: Option<Py<PyAny>> = shutdown_event.map(\|e\| e.unbind());

	future_into_py(py, async {
	let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();

	let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
	let handle_consume: JoinHandle<PyResult<Result<(), IggyError>>> =
	get_runtime().spawn(scope(task_locals, async move {
	let task_locals =
	Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
	let consumer = PyCallbackConsumer {
	callback: Arc::new(callback),
	task_locals: Arc::new(Mutex::new(task_locals)),
	};
	let mut inner = inner.lock().await;
	Ok(inner.consume_messages(&consumer, shutdown_rx).await)
	}));
	let consume_result;

	if let Some(shutdown_event) = shutdown_event {
	let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
	async fn shutdown_impl(
	shutdown_event: Py<PyAny>,
	shutdown_tx: Sender<()>,
	) -> PyResult<()> {
	Python::attach(\|py\| {
	into_future(shutdown_event.bind(py).as_any().call_method0("wait")?)
	})?
	.await?;
	shutdown_tx.send(()).map_err(\|_\| {
	PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
	"Failed to signal shutdown",
	)
	})?;
	Ok(())
	}
	let handle_shutdown: JoinHandle<Result<(), PyErr>> = get_runtime().spawn(scope(
	task_locals,
	shutdown_impl(shutdown_event, shutdown_tx),
	));
	let shutdown_result;
	(consume_result, shutdown_result) = tokio::join!(handle_consume, handle_shutdown);
	shutdown_result.map_err(\|e\| {
	PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
	})??;
	} else {
	consume_result = handle_consume.await;
	}

	consume_result
	.map_err(\|e\| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))??
	.map_err(\|e\| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
	Ok(())
	})
	}
	}

	#[gen_stub_pyclass]
	#[pyclass]
	pub struct ConsumerGroup {
	pub(crate) inner: RustConsumerGroup,
	}

	impl From<RustConsumerGroup> for ConsumerGroup {
	fn from(group: RustConsumerGroup) -> Self {
	Self { inner: group }
	}
	}

	#[gen_stub_pymethods]
	#[pymethods]
	impl ConsumerGroup {
	/// Gets the unique identifier (numeric) of the consumer group.
	#[getter]
	pub fn id(&self) -> u32 {
	self.inner.id
	}

	/// Gets the name of the consumer group.
	#[getter]
	pub fn name(&self) -> String {
	self.inner.name.to_string()
	}

	/// Gets the number of partitions the consumer group is consuming.
	#[getter]
	pub fn partitions_count(&self) -> u32 {
	self.inner.partitions_count
	}

	/// Gets the number of members in the consumer group.
	#[getter]
	pub fn members_count(&self) -> u32 {
	self.inner.members_count
	}
	}

	#[gen_stub_pyclass]
	#[pyclass]
	pub struct ConsumerGroupDetails {
	pub(crate) inner: RustConsumerGroupDetails,
	}

	impl From<RustConsumerGroupDetails> for ConsumerGroupDetails {
	fn from(group: RustConsumerGroupDetails) -> Self {
	Self { inner: group }
	}
	}

	#[gen_stub_pymethods]
	#[pymethods]
	impl ConsumerGroupDetails {
	/// Gets the unique identifier (numeric) of the consumer group.
	#[getter]
	pub fn id(&self) -> u32 {
	self.inner.id
	}

	/// Gets the name of the consumer group.
	#[getter]
	pub fn name(&self) -> String {
	self.inner.name.to_string()
	}

	/// Gets the number of partitions the consumer group is consuming.
	#[getter]
	pub fn partitions_count(&self) -> u32 {
	self.inner.partitions_count
	}

	/// Gets the number of members in the consumer group.
	#[getter]
	pub fn members_count(&self) -> u32 {
	self.inner.members_count
	}

	/// Gets the collection of members in the consumer group.
	#[getter]
	pub fn members(&self) -> Vec<ConsumerGroupMember> {
	self.inner
	.members
	.iter()
	.map(ConsumerGroupMember::from)
	.collect()
	}
	}

	#[gen_stub_pyclass]
	#[pyclass]
	pub struct ConsumerGroupMember {
	pub(crate) inner: RustConsumerGroupMember,
	}

	impl From<&RustConsumerGroupMember> for ConsumerGroupMember {
	fn from(member: &RustConsumerGroupMember) -> Self {
	Self {
	inner: RustConsumerGroupMember {
	id: member.id,
	partitions_count: member.partitions_count,
	partitions: member.partitions.clone(),
	},
	}
	}
	}

	#[gen_stub_pymethods]
	#[pymethods]
	impl ConsumerGroupMember {
	/// Gets the unique identifier (numeric) of the consumer group member.
	#[getter]
	pub fn id(&self) -> u32 {
	self.inner.id
	}

	/// Gets the number of partitions the consumer group member is consuming.
	#[getter]
	pub fn partitions_count(&self) -> u32 {
	self.inner.partitions_count
	}

	/// Gets the collection of partitions the consumer group member is consuming.
	#[getter]
	pub fn partitions(&self) -> Vec<u32> {
	self.inner.partitions.clone()
	}
	}

	#[pyclass]
	pub struct ReceiveMessageIterator {
	pub(crate) inner: Arc<Mutex<RustIggyConsumer>>,
	}

	#[pymethods]
	impl ReceiveMessageIterator {
	pub fn __anext__<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
	let inner = self.inner.clone();
	future_into_py(py, async move {
	let mut inner = inner.lock().await;
	if let Some(message) = inner.next().await {
	Ok(message
	.map(\|m\| ReceiveMessage {
	inner: m.message,
	partition_id: m.partition_id,
	})
	.map_err(\|e\| {
	PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
	})?)
	} else {
	Err(PyStopAsyncIteration::new_err("No more messages"))
	}
	})
	}

	pub fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
	slf
	}
	}

	struct PyCallbackConsumer {
	callback: Arc<Py<PyAny>>,
	task_locals: Arc<Mutex<TaskLocals>>,
	}

	impl MessageConsumer for PyCallbackConsumer {
	async fn consume(&self, received: ReceivedMessage) -> Result<(), IggyError> {
	let callback = self.callback.clone();
	let task_locals = self.task_locals.clone().lock_owned().await;
	let task_locals = task_locals.clone();
	let message = ReceiveMessage {
	inner: received.message,
	partition_id: received.partition_id,
	};
	get_runtime()
	.spawn(scope(task_locals, async move {
	Python::attach(\|py\| {
	let callback = callback.bind(py);
	let result = callback.as_any().call1((message,))?;
	into_future(result)
	})
	}))
	.await
	.map_err(\|_\| IggyError::CannotReadMessage)?
	.map_err(\|_\| IggyError::CannotReadMessage)?
	.await
	.map_err(\|_\| IggyError::CannotReadMessage)?;
	Ok(())
	}
	}

	/// The auto-commit configuration for storing the offset on the server.
	// #[derive(Debug, PartialEq, Copy, Clone)]
	#[gen_stub_pyclass_complex_enum]
	#[pyclass]
	pub enum AutoCommit {
	/// The auto-commit is disabled and the offset must be stored manually by the consumer.
	Disabled(),
	/// The auto-commit is enabled and the offset is stored on the server after a certain interval.
	Interval(Py<PyDelta>),
	/// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode when consuming the messages.
	IntervalOrWhen(Py<PyDelta>, AutoCommitWhen),
	/// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode after consuming the messages.
	IntervalOrAfter(Py<PyDelta>, AutoCommitAfter),
	/// The auto-commit is enabled and the offset is stored on the server depending on the mode when consuming the messages.
	When(AutoCommitWhen),
	/// The auto-commit is enabled and the offset is stored on the server depending on the mode after consuming the messages.
	After(AutoCommitAfter),
	}

	impl From<&AutoCommit> for RustAutoCommit {
	fn from(val: &AutoCommit) -> RustAutoCommit {
	match val {
	AutoCommit::Disabled() => RustAutoCommit::Disabled,
	AutoCommit::Interval(delta) => {
	let duration = py_delta_to_iggy_duration(delta);
	RustAutoCommit::Interval(duration)
	}
	AutoCommit::IntervalOrWhen(delta, when) => {
	let duration = py_delta_to_iggy_duration(delta);
	RustAutoCommit::IntervalOrWhen(duration, when.into())
	}
	AutoCommit::IntervalOrAfter(delta, after) => {
	let duration = py_delta_to_iggy_duration(delta);
	RustAutoCommit::IntervalOrAfter(duration, after.into())
	}
	AutoCommit::When(when) => RustAutoCommit::When(when.into()),
	AutoCommit::After(after) => RustAutoCommit::After(after.into()),
	}
	}
	}

	/// The auto-commit mode for storing the offset on the server.
	#[derive(Debug, PartialEq, Copy, Clone)]
	#[gen_stub_pyclass_complex_enum(skip_stub_type)]
	#[pyclass(from_py_object)]
	pub enum AutoCommitWhen {
	/// The offset is stored on the server when the messages are received.
	PollingMessages(),
	/// The offset is stored on the server when all the messages are consumed.
	ConsumingAllMessages(),
	/// The offset is stored on the server when consuming each message.
	ConsumingEachMessage(),
	/// The offset is stored on the server when consuming every Nth message.
	ConsumingEveryNthMessage(u32),
	}

	impl From<&AutoCommitWhen> for RustAutoCommitWhen {
	fn from(val: &AutoCommitWhen) -> RustAutoCommitWhen {
	match val {
	AutoCommitWhen::PollingMessages() => RustAutoCommitWhen::PollingMessages,
	AutoCommitWhen::ConsumingAllMessages() => RustAutoCommitWhen::ConsumingAllMessages,
	AutoCommitWhen::ConsumingEachMessage() => RustAutoCommitWhen::ConsumingEachMessage,
	AutoCommitWhen::ConsumingEveryNthMessage(n) => {
	RustAutoCommitWhen::ConsumingEveryNthMessage(n.to_owned())
	}
	}
	}
	}

	impl PyStubType for AutoCommitWhen {
	fn type_output() -> TypeInfo {
	TypeInfo::unqualified("AutoCommitWhen")
	}
	}

	/// The auto-commit mode for storing the offset on the server after receiving the messages.
	#[derive(Debug, PartialEq, Copy, Clone)]
	#[gen_stub_pyclass_complex_enum(skip_stub_type)]
	#[pyclass(from_py_object)]
	#[allow(clippy::enum_variant_names)]
	pub enum AutoCommitAfter {
	/// The offset is stored on the server after all the messages are consumed.
	ConsumingAllMessages(),
	/// The offset is stored on the server after consuming each message.
	ConsumingEachMessage(),
	/// The offset is stored on the server after consuming every Nth message.
	ConsumingEveryNthMessage(u32),
	}

	impl From<&AutoCommitAfter> for RustAutoCommitAfter {
	fn from(val: &AutoCommitAfter) -> RustAutoCommitAfter {
	match val {
	AutoCommitAfter::ConsumingAllMessages() => RustAutoCommitAfter::ConsumingAllMessages,
	AutoCommitAfter::ConsumingEachMessage() => RustAutoCommitAfter::ConsumingEachMessage,
	AutoCommitAfter::ConsumingEveryNthMessage(n) => {
	RustAutoCommitAfter::ConsumingEveryNthMessage(n.to_owned())
	}
	}
	}
	}

	impl PyStubType for AutoCommitAfter {
	fn type_output() -> TypeInfo {
	TypeInfo::unqualified("AutoCommitAfter")
	}
	}

	pub fn py_delta_to_iggy_duration(delta1: &Py<PyDelta>) -> IggyDuration {
	Python::attach(\|py\| {
	let delta = delta1.bind(py);
	let seconds = (delta.get_days() * 60 * 60 * 24 + delta.get_seconds()) as u64;
	let nanos = (delta.get_microseconds() * 1_000) as u32;
	IggyDuration::new(Duration::new(seconds, nanos))
	})
	}