Address feedback on CallState params and adding more tests

Author: k-raina

src/confluent_kafka/src/Admin.c

@@ -4717,7 +4717,7 @@ static void Admin_background_event_cb (rd_kafka_t *rk, rd_kafka_event_t *rkev,
         PyGILState_STATE gstate;
         PyObject *error, *method, *ret;
         PyObject *result = NULL;
-        PyObject *exctype = NULL, *exc = NULL, *excargs = NULL;
+        PyObject *exc = NULL, *excargs = NULL;
 
         /* Acquire GIL */
         gstate = PyGILState_Ensure();
@@ -5093,7 +5093,7 @@ static void Admin_background_event_cb (rd_kafka_t *rk, rd_kafka_event_t *rkev,
                         PyObject *trace = NULL;
 
                         /* Fetch (and clear) currently raised exception */
-                        cfl_exception_fetch(&exctype, &error, &trace);
+                        cfl_exception_fetch(&exc);
                         Py_XDECREF(trace);
                 }
 
@@ -5124,22 +5124,17 @@ static void Admin_background_event_cb (rd_kafka_t *rk, rd_kafka_event_t *rkev,
          * Pass an exception to future.set_exception().
          */
 
-        if (!exctype) {
+        if (!exc) {
                 /* No previous exception raised, use KafkaException */
-                exctype = KafkaException;
-                Py_INCREF(exctype);
-        }
-
-        /* Create a new exception based on exception type and error. */
-        excargs = PyTuple_New(1);
-        Py_INCREF(error); /* tuple's reference */
-        PyTuple_SET_ITEM(excargs, 0, error);
-        exc = ((PyTypeObject *)exctype)->tp_new(
-                (PyTypeObject *)exctype, NULL, NULL);
-        exc->ob_type->tp_init(exc, excargs, NULL);
-        Py_DECREF(excargs);
-        Py_XDECREF(exctype);
-        Py_XDECREF(error); /* from error source above */
+                excargs = PyTuple_New(1);
+                Py_INCREF(error); /* tuple's reference */
+                PyTuple_SET_ITEM(excargs, 0, error);
+                exc = ((PyTypeObject *)KafkaException)->tp_new(
+                        (PyTypeObject *)KafkaException, NULL, NULL);
+                exc->ob_type->tp_init(exc, excargs, NULL);
+                Py_DECREF(excargs);
+                Py_XDECREF(error); /* from error source above */
+        }
 
         /*
          * Call future.set_exception(exc)

src/confluent_kafka/src/Consumer.c

@@ -424,6 +424,7 @@ static void Consumer_offset_commit_cb (rd_kafka_t *rk, rd_kafka_resp_err_t err,
         if (result)
                 Py_DECREF(result);
         else {
+                CallState_fetch_exception(cs);
                 CallState_crash(cs);
                 rd_kafka_yield(rk);
         }
@@ -1586,7 +1587,6 @@ static void Consumer_rebalance_cb (rd_kafka_t *rk, rd_kafka_resp_err_t err,
 		if (result)
 			Py_DECREF(result);
 		else {
-			
 			CallState_fetch_exception(cs);
 			CallState_crash(cs);
 			rd_kafka_yield(rk);

src/confluent_kafka/src/Producer.c

@@ -174,7 +174,6 @@ static void dr_msg_cb (rd_kafka_t *rk, const rd_kafka_message_t *rkm,
 	if (result)
 		Py_DECREF(result);
 	else {
-		
 		CallState_fetch_exception(cs);
 		CallState_crash(cs);
 		rd_kafka_yield(rk);

src/confluent_kafka/src/confluent_kafka.c

@@ -1766,8 +1766,7 @@ static void error_cb (rd_kafka_t *rk, int err, const char *reason, void *opaque)
 	if (result)
 		Py_DECREF(result);
 	else {
-        
-        CallState_fetch_exception(cs);
+		CallState_fetch_exception(cs);
         crash:
 		CallState_crash(cs);
 		rd_kafka_yield(h->rk);
@@ -2589,9 +2588,7 @@ void CallState_begin (Handle *h, CallState *cs) {
 	cs->thread_state = PyEval_SaveThread();
 	assert(cs->thread_state != NULL);
 	cs->crashed = 0;
-	cs->exception_type = NULL;
 	cs->exception_value = NULL;
-	cs->exception_traceback = NULL;
 #ifdef WITH_PY_TSS
         PyThread_tss_set(&h->tlskey, cs);
 #else
@@ -2617,7 +2614,7 @@ int CallState_end (Handle *h, CallState *cs) {
 
 	if (cs->crashed) {
 		/* Restore the saved exception if we have one */
-		if (cs->exception_type) {
+		if (cs->exception_value) {
 			CallState_restore_exception(cs);
 		}
 		return 0;

src/confluent_kafka/src/confluent_kafka.h

@@ -270,9 +270,7 @@ int  Handle_traverse (Handle *h, visitproc visit, void *arg);
 typedef struct {
 	PyThreadState *thread_state;
 	int crashed;   /* Callback crashed */
-	PyObject *exception_type;    /* Stored exception type */
 	PyObject *exception_value;   /* Stored exception value */
-	PyObject *exception_traceback; /* Stored exception traceback */
 } CallState;
 
 /**
@@ -282,51 +280,42 @@ typedef struct {
  */
 
 static inline void
-cfl_exception_fetch(PyObject **exc_type, PyObject **exc_value, PyObject **exc_traceback) {
+cfl_exception_fetch(PyObject **exc_value) {
 #if PY_VERSION_HEX >= 0x030c0000
     /* Python 3.12+ - use new API */
-    PyObject *exc = PyErr_GetRaisedException();
-    if (exc) {
-        *exc_type = (PyObject *)Py_TYPE(exc);
-        Py_INCREF(*exc_type);
-        *exc_value = exc;
-        *exc_traceback = PyException_GetTraceback(exc);
-    } else {
-        *exc_type = *exc_value = *exc_traceback = NULL;
-    }
+    *exc_value = PyErr_GetRaisedException();
 #else
     /* Python < 3.12 - use legacy API */
-    PyErr_Fetch(exc_type, exc_value, exc_traceback);
+    PyObject *exc_type, *exc_traceback;
+    PyErr_Fetch(&exc_type, exc_value, &exc_traceback);
+    Py_XDECREF(exc_type);
+    Py_XDECREF(exc_traceback);
 #endif
 }
 
 static inline void
-cfl_exception_restore(PyObject *exc_type, PyObject *exc_value, PyObject *exc_traceback) {
+cfl_exception_restore(PyObject *exc_value) {
 #if PY_VERSION_HEX >= 0x030c0000
     /* Python 3.12+ - use new API */
     if (exc_value) {
         PyErr_SetRaisedException(exc_value);
-        Py_XDECREF(exc_type);
-        Py_XDECREF(exc_traceback);
     }
 #else
     /* Python < 3.12 - use legacy API */
-    PyErr_Restore(exc_type, exc_value, exc_traceback);
+    PyErr_SetObject(PyExceptionInstance_Class(exc_value), exc_value);
 #endif
 }
 
 static inline void
 CallState_fetch_exception(CallState *cs) {
-    cfl_exception_fetch(&cs->exception_type, &cs->exception_value, &cs->exception_traceback);
+    cfl_exception_fetch(&cs->exception_value);
 }
 
 static inline void
 CallState_restore_exception(CallState *cs) {
-    if (cs->exception_type) {
-        cfl_exception_restore(cs->exception_type, cs->exception_value, cs->exception_traceback);
-        cs->exception_type = NULL;
+    if (cs->exception_value) {
+        cfl_exception_restore(cs->exception_value);
         cs->exception_value = NULL;
-        cs->exception_traceback = NULL;
     }
 }
 

tests/test_Admin.py

@@ -1224,3 +1224,89 @@ def test_elect_leaders():
     with pytest.raises(KafkaException):
         a.elect_leaders(correct_election_type, [correct_partitions])\
             .result(timeout=1)
+
+
+@pytest.mark.skipif(libversion()[1] < 0x000b0500,
+                    reason="AdminAPI requires librdkafka >= v0.11.5")
+def test_admin_callback_exception_no_system_error():
+    """Test AdminClient callbacks exception handling with different exception types"""
+
+    # Test error_cb with different exception types
+    def error_cb_kafka_exception(error):
+        raise KafkaException(KafkaError._FAIL, "KafkaException from error_cb")
+
+    def error_cb_value_error(error):
+        raise ValueError("ValueError from error_cb")
+
+    def error_cb_runtime_error(error):
+        raise RuntimeError("RuntimeError from error_cb")
+
+    # Test error_cb with KafkaException
+    admin = AdminClient({
+        'bootstrap.servers': 'nonexistent-broker:9092',
+        'socket.timeout.ms': 100,
+        'error_cb': error_cb_kafka_exception
+    })
+
+    with pytest.raises(KafkaException) as exc_info:
+        admin.poll(timeout=0.2)
+    assert "KafkaException from error_cb" in str(exc_info.value)
+
+    # Test error_cb with ValueError
+    admin = AdminClient({
+        'bootstrap.servers': 'nonexistent-broker:9092',
+        'socket.timeout.ms': 100,
+        'error_cb': error_cb_value_error
+    })
+
+    with pytest.raises(ValueError) as exc_info:
+        admin.poll(timeout=0.2)
+    assert "ValueError from error_cb" in str(exc_info.value)
+
+    # Test error_cb with RuntimeError
+    admin = AdminClient({
+        'bootstrap.servers': 'nonexistent-broker:9092',
+        'socket.timeout.ms': 100,
+        'error_cb': error_cb_runtime_error
+    })
+
+    with pytest.raises(RuntimeError) as exc_info:
+        admin.poll(timeout=0.2)
+    assert "RuntimeError from error_cb" in str(exc_info.value)
+
+
+@pytest.mark.skipif(libversion()[1] < 0x000b0500,
+                    reason="AdminAPI requires librdkafka >= v0.11.5")
+def test_admin_multiple_callbacks_different_error_types():
+    """Test AdminClient with multiple callbacks configured with different error types to see which one gets triggered"""
+
+    callbacks_called = []
+
+    def error_cb_that_raises_runtime(error):
+        callbacks_called.append('error_cb_runtime')
+        raise RuntimeError("RuntimeError from error_cb")
+
+    def stats_cb_that_raises_value(stats_json):
+        callbacks_called.append('stats_cb_value')
+        raise ValueError("ValueError from stats_cb")
+
+    def throttle_cb_that_raises_kafka(throttle_event):
+        callbacks_called.append('throttle_cb_kafka')
+        raise KafkaException(KafkaError._FAIL, "KafkaException from throttle_cb")
+
+    admin = AdminClient({
+        'bootstrap.servers': 'nonexistent-broker:9092',
+        'socket.timeout.ms': 100,
+        'statistics.interval.ms': 100,  # Enable stats callback
+        'error_cb': error_cb_that_raises_runtime,
+        'stats_cb': stats_cb_that_raises_value,
+        'throttle_cb': throttle_cb_that_raises_kafka
+    })
+
+    # Test that error_cb callback raises an exception (it's triggered by connection failures)
+    with pytest.raises(RuntimeError) as exc_info:
+        admin.poll(timeout=0.2)
+
+    # Verify that error_cb was called
+    assert len(callbacks_called) > 0
+    assert 'error_cb_runtime' in callbacks_called