Address feedback

Author: k-raina

src/confluent_kafka/cimpl.pyi

@@ -558,8 +558,10 @@ class ShareConsumer:
     def subscribe(self, topics: List[str]) -> None: ...
     def unsubscribe(self) -> None: ...
     def subscription(self) -> List[str]: ...
-    def consume_batch(self, timeout: float = -1) -> List[Message]: ...
+    def poll(self, timeout: float = -1) -> List[Message]: ...
     def close(self) -> None: ...
+    def __enter__(self) -> "ShareConsumer": ...
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_traceback: Any) -> Optional[bool]: ...
 
 class _AdminClientImpl:
     def __init__(self, config: Dict[str, Union[str, int, float, bool]]) -> None: ...

src/confluent_kafka/src/ShareConsumer.c

@@ -61,6 +61,8 @@ static void ShareConsumer_dealloc(ShareConsumerHandle *self) {
         if (self->rkshare) {
                 CallState cs;
                 CallState_begin((Handle *)self, &cs);
+                /* TODO KIP-932: Use rd_kafka_share_destroy_flags() once
+                 * available in the librdkafka public API. */
                 rd_kafka_share_destroy(self->rkshare);
                 self->rkshare = NULL;
                 CallState_end((Handle *)self, &cs);
@@ -197,12 +199,12 @@ static PyObject *ShareConsumer_subscription(ShareConsumerHandle *self) {
 
 
 /**
- * @brief Consume a batch of messages from the share consumer.
+ * @brief Poll for a batch of messages from the share consumer.
  *
  */
-static PyObject *ShareConsumer_consume_batch(ShareConsumerHandle *self,
-                                             PyObject *args,
-                                             PyObject *kwargs) {
+static PyObject *ShareConsumer_poll(ShareConsumerHandle *self,
+                                    PyObject *args,
+                                    PyObject *kwargs) {
         double tmout                    = -1.0f;
         static char *kws[]              = {"timeout", NULL};
         rd_kafka_message_t **rkmessages = NULL;
@@ -299,20 +301,7 @@ static PyObject *ShareConsumer_consume_batch(ShareConsumerHandle *self,
 
         /* Handle error from rd_kafka_share_consume_batch() */
         if (error) {
-                const char *error_str = rd_kafka_error_string(error);
-                int is_fatal          = rd_kafka_error_is_fatal(error);
-                int is_retriable      = rd_kafka_error_is_retriable(error);
-
-                if (is_fatal) {
-                        PyErr_Format(PyExc_RuntimeError, "Fatal error: %s",
-                                     error_str);
-                } else {
-                        PyErr_Format(KafkaException,
-                                     "Error: %s (retriable: %s)", error_str,
-                                     is_retriable ? "yes" : "no");
-                }
-
-                rd_kafka_error_destroy(error);
+                cfl_PyErr_from_error_destroy(error);
                 free(rkmessages);
                 return NULL;
         }
@@ -348,6 +337,8 @@ static PyObject *ShareConsumer_close(ShareConsumerHandle *self) {
                 Py_RETURN_NONE;
 
         CallState_begin((Handle *)self, &cs);
+        /* TODO KIP-932: rd_kafka_share_consumer_close() return type will change
+         * to rd_kafka_error_t *. Update error handling accordingly. */
         err = rd_kafka_share_consumer_close(self->rkshare);
         rd_kafka_share_destroy(self->rkshare);
         self->rkshare = NULL;
@@ -364,6 +355,36 @@ static PyObject *ShareConsumer_close(ShareConsumerHandle *self) {
 }
 
 
+/**
+ * @brief Context manager entry — returns self.
+ */
+static PyObject *ShareConsumer_enter(ShareConsumerHandle *self) {
+        Py_INCREF(self);
+        return (PyObject *)self;
+}
+
+/**
+ * @brief Context manager exit — calls close().
+ */
+static PyObject *ShareConsumer_exit(ShareConsumerHandle *self, PyObject *args) {
+        PyObject *exc_type, *exc_value, *exc_traceback;
+
+        if (!PyArg_UnpackTuple(args, "__exit__", 3, 3, &exc_type, &exc_value,
+                               &exc_traceback))
+                return NULL;
+
+        /* Cleanup: call close() */
+        if (self->rkshare) {
+                PyObject *result = ShareConsumer_close(self);
+                if (!result)
+                        return NULL;
+                Py_DECREF(result);
+        }
+
+        Py_RETURN_NONE;
+}
+
+
 /**
  * @brief ShareConsumer methods.
  */
@@ -399,14 +420,10 @@ static PyMethodDef ShareConsumer_methods[] = {
      "  :raises RuntimeError: if called on a closed share consumer\n"
      "\n"},
 
-    {"consume_batch", (PyCFunction)ShareConsumer_consume_batch,
-     METH_VARARGS | METH_KEYWORDS,
-     ".. py:function:: consume_batch([timeout=-1])\n"
+    {"poll", (PyCFunction)ShareConsumer_poll, METH_VARARGS | METH_KEYWORDS,
+     ".. py:function:: poll([timeout=-1])\n"
      "\n"
-     "  Consume a batch of messages from the share consumer.\n"
-     "\n"
-     "  This is the ONLY consumption method for ShareConsumer.\n"
-     "  Share consumers do NOT have a poll() method - they are batch-only.\n"
+     "  Poll for a batch of messages from the share consumer.\n"
      "\n"
      "  The application must check each Message object's error() method\n"
      "  to distinguish between proper messages (error() returns None)\n"
@@ -421,9 +438,8 @@ static PyMethodDef ShareConsumer_methods[] = {
      "                        Default: -1 (infinite)\n"
      "  :returns: List of Message objects (possibly empty on timeout)\n"
      "  :rtype: list(Message)\n"
-     "  :raises RuntimeError: if called on a closed share consumer or on "
-     "fatal error\n"
-     "  :raises KafkaException: on non-fatal errors\n"
+     "  :raises KafkaException: on error\n"
+     "  :raises RuntimeError: if called on a closed share consumer\n"
      "  :raises KeyboardInterrupt: if Ctrl+C pressed during consumption\n"
      "\n"},
 
@@ -438,6 +454,15 @@ static PyMethodDef ShareConsumer_methods[] = {
      "  :raises KafkaException: on error\n"
      "\n"},
 
+    /* TODO KIP-932: Add set_sasl_credentials once librdkafka exposes
+     * rd_kafka_sasl_set_credentials() (or the underlying rd_kafka_t *)
+     * for rd_kafka_share_t handles. */
+
+    {"__enter__", (PyCFunction)ShareConsumer_enter, METH_NOARGS,
+     "Context manager entry."},
+    {"__exit__", (PyCFunction)ShareConsumer_exit, METH_VARARGS,
+     "Context manager exit. Automatically closes the share consumer."},
+
     {NULL}};
 
 
@@ -530,10 +555,6 @@ PyTypeObject ShareConsumerType = {
     "assigned to multiple consumers. Messages are delivered to only one "
     "consumer.\n"
     "\n"
-    ".. note::\n"
-    "   ShareConsumer only supports batch consumption via consume_batch().\n"
-    "   There is NO poll() method for single messages.\n"
-    "\n"
     ":param dict config: Configuration properties. At a minimum, "
     "``group.id`` **must** be set and ``bootstrap.servers`` **should** be "
     "set.\n"

tests/test_ShareConsumer.py

@@ -66,8 +66,8 @@ def test_unsubscribe():
     sc.close()
 
 
-def test_consume_batch_no_broker():
-    """Test consume_batch() returns empty list when no broker available."""
+def test_poll_no_broker():
+    """Test poll() returns empty list when no broker available."""
     sc = ShareConsumer(
         {
             'group.id': 'test-share-group',
@@ -79,13 +79,33 @@ def test_consume_batch_no_broker():
     sc.subscribe(['test-topic'])
 
     # Should timeout and return empty list
-    messages = sc.consume_batch(timeout=0.1)
+    messages = sc.poll(timeout=0.1)
     assert isinstance(messages, list)
     # May be empty or contain error messages
 
     sc.close()
 
 
+def test_context_manager():
+    """Test that ShareConsumer works as a context manager and closes on exit."""
+    with ShareConsumer(
+        {
+            'group.id': 'test-share-group',
+            'bootstrap.servers': 'localhost:9092',
+            'socket.timeout.ms': 100,
+        }
+    ) as sc:
+        assert sc is not None
+        sc.subscribe(['test-topic'])
+        subscription = sc.subscription()
+        assert 'test-topic' in subscription
+
+    # After exiting the context manager, the consumer should be closed
+    with pytest.raises(RuntimeError) as ex:
+        sc.subscribe(['test-topic'])
+    assert ex.match('Share consumer closed')
+
+
 def test_close_idempotent():
     """Test that close() can be called multiple times."""
     sc = ShareConsumer(
@@ -126,7 +146,7 @@ def test_any_method_after_close_throws_exception():
     assert ex.match('Share consumer closed')
 
     with pytest.raises(RuntimeError) as ex:
-        sc.consume_batch(timeout=0.1)
+        sc.poll(timeout=0.1)
     assert ex.match('Share consumer closed')
 
 
@@ -159,8 +179,8 @@ def test_concurrent_consumers():
         sc1.subscribe(['test-topic'])
         sc2.subscribe(['test-topic'])
 
-        messages1 = sc1.consume_batch(timeout=2.0)
-        messages2 = sc2.consume_batch(timeout=2.0)
+        messages1 = sc1.poll(timeout=2.0)
+        messages2 = sc2.poll(timeout=2.0)
 
         # Verify no overlap (share group semantics)
         offsets1 = {(msg.topic(), msg.partition(), msg.offset()) for msg in messages1 if not msg.error()}
@@ -190,14 +210,14 @@ def my_error_cb(error):
     )
 
     sc.subscribe(['test-topic'])
-    sc.consume_batch(timeout=0.5)
+    sc.poll(timeout=0.5)
 
     assert len(error_called) > 0, "error_cb should have been called"
     sc.close()
 
 
 def test_error_cb_exception_propagates():
-    """Test that an exception raised in error_cb propagates to consume_batch."""
+    """Test that an exception raised in error_cb propagates to poll."""
     error_called = []
 
     def error_cb_that_raises(error):
@@ -216,7 +236,7 @@ def error_cb_that_raises(error):
     sc.subscribe(['test-topic'])
 
     with pytest.raises(RuntimeError) as exc_info:
-        sc.consume_batch(timeout=0.5)
+        sc.poll(timeout=0.5)
 
     assert "Test exception from error_cb" in str(exc_info.value)
     assert len(error_called) > 0
@@ -245,7 +265,7 @@ def my_stats_cb(stats_json):
     )
 
     sc.subscribe(['test-topic'])
-    sc.consume_batch(timeout=0.5)
+    sc.poll(timeout=0.5)
 
     assert len(stats_called) > 0, "stats_cb should have been called"
     # Verify we got valid JSON string
@@ -257,7 +277,7 @@ def my_stats_cb(stats_json):
 
 
 def test_stats_cb_exception_propagates():
-    """Test that an exception raised in stats_cb propagates to consume_batch."""
+    """Test that an exception raised in stats_cb propagates to poll."""
     stats_called = []
 
     def stats_cb_that_raises(stats_json):
@@ -277,7 +297,7 @@ def stats_cb_that_raises(stats_json):
     sc.subscribe(['test-topic'])
 
     with pytest.raises(RuntimeError) as exc_info:
-        sc.consume_batch(timeout=0.5)
+        sc.poll(timeout=0.5)
 
     assert "Test exception from stats_cb" in str(exc_info.value)
     assert len(stats_called) > 0
@@ -307,5 +327,5 @@ def my_throttle_cb(event):
     sc.subscribe(['test-topic'])
 
     # throttle_cb won't fire without broker throttling — just verify no crash
-    sc.consume_batch(timeout=0.2)
+    sc.poll(timeout=0.2)
     sc.close()