Add authentication credential callback support

docs/api_reference.rst

@@ -187,6 +187,8 @@ Testing Utilities
 Helper Types
 ============
 
+.. autoclass:: blazingmq.AuthnCredentialProvider
+
 .. autoclass:: blazingmq.PropertyTypeDict
 
 .. autoclass:: blazingmq.PropertyValueDict

docs/conf.py

@@ -185,6 +185,9 @@
 
 autodoc_typehints = 'description'
 autodoc_typehints_description_target = 'documented'
+autodoc_type_aliases = {
+    "AuthnCredentialProvider": "blazingmq.AuthnCredentialProvider",
+}
 
 
 def document_as_keyword_only(app, what, name, obj, options, signature, return_hint):

src/blazingmq/__init__.py

@@ -27,13 +27,15 @@
 from ._session import Session
 from ._session import SessionOptions
 from ._timeouts import Timeouts
+from ._typing import AuthnCredentialProvider
 from ._typing import PropertyTypeDict
 from ._typing import PropertyValueDict
 from .exceptions import Error
 
 __all__ = [
     "Ack",
     "AckStatus",
+    "AuthnCredentialProvider",
     "BasicHealthMonitor",
     "CompressionAlgorithmType",
     "Error",

src/blazingmq/_ext.pyi

@@ -37,6 +37,9 @@ class FakeHostHealthMonitor:
     def set_healthy(self) -> None: ...
     def set_unhealthy(self) -> None: ...
 
+class AuthnCredentialCbAdapter:
+    def __init__(self, callback: Callable[[], Optional[tuple[str, bytes]]]) -> None: ...
+
 class Session:
     def __init__(
         self,
@@ -53,6 +56,7 @@ class Session:
         timeouts: Timeouts = Timeouts(),
         monitor_host_health: bool = False,
         fake_host_health_monitor: Optional[FakeHostHealthMonitor] = None,
+        authn_credential_cb: Optional[AuthnCredentialCbAdapter] = None,
     ) -> None: ...
     def stop(self) -> None: ...
     def open_queue_sync(

src/blazingmq/_ext.pyx

@@ -21,12 +21,15 @@ import weakref
 from bsl cimport optional
 from bsl cimport pair
 from bsl cimport shared_ptr
+from bsl cimport string
+from bsl cimport vector
 from bsl.bsls cimport TimeInterval
 from cpython.ceval cimport PyEval_InitThreads
 from libcpp cimport bool as cppbool
 
 from bmq.bmqa cimport ManualHostHealthMonitor
 from bmq.bmqt cimport AckResult
+from bmq.bmqt cimport AuthnCredential
 from bmq.bmqt cimport CompressionAlgorithmType
 from bmq.bmqt cimport HostHealthState
 from bmq.bmqt cimport PropertyType
@@ -153,6 +156,36 @@ cdef class FakeHostHealthMonitor:
             self._monitor.get().setState(HostHealthState.e_UNHEALTHY)
 
 
+cdef class AuthnCredentialCbAdapter:
+    cdef object _callback  # Store the Python callable
+
+    def __cinit__(self, callback):
+        self._callback = callback
+
+    # This method will be called by C++ code via PyObject_CallMethod
+    # Returns None for no credential, or (mechanism, data) tuple
+    def get_credential_data(self):
+        try:
+            result = self._callback()
+            if result is None:
+                return None
+
+            if not isinstance(result, tuple) or len(result) != 2:
+                raise ValueError("callback must return (str, bytes) or None")
+
+            mechanism, data = result
+            if not isinstance(mechanism, str) or not isinstance(data, bytes):
+                raise ValueError("callback must return (str, bytes) or None")
+
+            # Return as-is, let C++ side handle conversion
+            return result
+
+        except Exception:
+            # Log error or handle as needed
+            LOGGER.exception("Error in authentication credential callback")
+            return None
+
+
 cdef class Session:
     cdef object __weakref__
     cdef NativeSession* _session
@@ -174,6 +207,7 @@ cdef class Session:
         timeouts: _timeouts.Timeouts = _timeouts.Timeouts(),
         monitor_host_health: bool = False,
         fake_host_health_monitor: FakeHostHealthMonitor = None,
+        authn_credential_cb: AuthnCredentialCbAdapter = None,
         _mock: Optional[object] = None,
     ) -> None:
         cdef shared_ptr[ManualHostHealthMonitor] fake_host_health_monitor_sp
@@ -225,6 +259,7 @@ cdef class Session:
             session_cb,
             message_cb,
             ack_cb,
+            authn_credential_cb,
             c_broker_uri,
             c_script_name,
             COMPRESSION_ALGO_FROM_PY_MAPPING[message_compression_algorithm],

src/blazingmq/_session.py

@@ -25,6 +25,7 @@
 from . import _six as six
 from ._enums import CompressionAlgorithmType
 from ._enums import PropertyType
+from ._ext import AuthnCredentialCbAdapter
 from ._ext import DEFAULT_CONSUMER_PRIORITY
 from ._ext import DEFAULT_MAX_UNCONFIRMED_BYTES
 from ._ext import DEFAULT_MAX_UNCONFIRMED_MESSAGES
@@ -36,6 +37,7 @@
 from ._messages import MessageHandle
 from ._monitors import BasicHealthMonitor
 from ._timeouts import Timeouts
+from ._typing import AuthnCredentialProvider
 from ._typing import PropertyTypeDict
 from ._typing import PropertyValueDict
 from ._typing import PropertyValueType
@@ -52,6 +54,10 @@ def DefaultMonitor() -> Union[BasicHealthMonitor, None]:
     return None
 
 
+def DefaultAuthnCredentialProvider() -> Optional[AuthnCredentialProvider]:
+    return None
+
+
 DEFAULT_TIMEOUT = DefaultTimeoutType()
 KNOWN_MONITORS = ("blazingmq.BasicHealthMonitor",)
 
@@ -265,6 +271,11 @@ class SessionOptions:
             healthy, `.HostUnhealthy` and `.HostHealthRestored` events with
             never be emitted, and the *suspends_on_bad_host_health* option of
             `QueueOptions` cannot be used.
+        authn_credential_provider:
+            An optional callable that returns authentication credentials as a
+            ``(mechanism, data)`` tuple of ``(str, bytes)``, or ``None`` if no
+            credentials are available.  If not provided, no authentication
+            credentials are sent to the broker.
         num_processing_threads:
             The number of threads for the SDK to use for processing events.
             This defaults to 1.
@@ -295,6 +306,9 @@ def __init__(
         message_compression_algorithm: Optional[CompressionAlgorithmType] = None,
         timeouts: Optional[Timeouts] = None,
         host_health_monitor: Union[BasicHealthMonitor, None] = (DefaultMonitor()),
+        authn_credential_provider: Optional[AuthnCredentialProvider] = (
+            DefaultAuthnCredentialProvider()
+        ),
         num_processing_threads: Optional[int] = None,
         blob_buffer_size: Optional[int] = None,
         channel_high_watermark: Optional[int] = None,
@@ -304,6 +318,7 @@ def __init__(
         self.message_compression_algorithm = message_compression_algorithm
         self.timeouts = timeouts
         self.host_health_monitor = host_health_monitor
+        self.authn_credential_provider = authn_credential_provider
         self.num_processing_threads = num_processing_threads
         self.blob_buffer_size = blob_buffer_size
         self.channel_high_watermark = channel_high_watermark
@@ -317,6 +332,7 @@ def __eq__(self, other: object) -> bool:
             self.message_compression_algorithm == other.message_compression_algorithm
             and self.timeouts == other.timeouts
             and self.host_health_monitor == other.host_health_monitor
+            and self.authn_credential_provider == other.authn_credential_provider
             and self.num_processing_threads == other.num_processing_threads
             and self.blob_buffer_size == other.blob_buffer_size
             and self.channel_high_watermark == other.channel_high_watermark
@@ -332,6 +348,7 @@ def __repr__(self) -> str:
             "message_compression_algorithm",
             "timeouts",
             "host_health_monitor",
+            "authn_credential_provider",
             "num_processing_threads",
             "blob_buffer_size",
             "channel_high_watermark",
@@ -379,6 +396,10 @@ class Session:
             `.HostHealthRestored` events will never be emitted, and the
             *suspends_on_bad_host_health* option of `QueueOptions` cannot be
             used.
+        authn_credential_provider: an optional callable that returns authentication
+            credentials as a ``(mechanism, data)`` tuple of ``(str, bytes)``,
+            or ``None`` if no credentials are available.  If not provided, no
+            authentication credentials are sent to the broker.
         num_processing_threads: The number of threads for the SDK to use for
             processing events.  This defaults to 1.
         blob_buffer_size: The size (in bytes) of the blob buffers to use.  This
@@ -418,6 +439,9 @@ def __init__(
         ),
         timeout: Union[Timeouts, float] = DEFAULT_TIMEOUT,
         host_health_monitor: Union[BasicHealthMonitor, None] = (DefaultMonitor()),
+        authn_credential_provider: Optional[AuthnCredentialProvider] = (
+            DefaultAuthnCredentialProvider()
+        ),
         num_processing_threads: Optional[int] = None,
         blob_buffer_size: Optional[int] = None,
         channel_high_watermark: Optional[int] = None,
@@ -433,6 +457,11 @@ def __init__(
 
         monitor_host_health = host_health_monitor is not None
         fake_host_health_monitor = getattr(host_health_monitor, "_monitor", None)
+        authn_credential_cb = (
+            AuthnCredentialCbAdapter(authn_credential_provider)
+            if authn_credential_provider is not None
+            else None
+        )
 
         self._has_no_on_message = on_message is None
 
@@ -459,6 +488,7 @@ def __init__(
             timeouts=_validate_timeouts(timeout),
             monitor_host_health=monitor_host_health,
             fake_host_health_monitor=fake_host_health_monitor,
+            authn_credential_cb=authn_credential_cb,
         )
         self._ext.set_owned_by_session()
 
@@ -511,6 +541,7 @@ def with_options(
                 message_compression_algorithm,
                 DEFAULT_TIMEOUT,
                 session_options.host_health_monitor,
+                session_options.authn_credential_provider,
                 session_options.num_processing_threads,
                 session_options.blob_buffer_size,
                 session_options.channel_high_watermark,
@@ -525,6 +556,7 @@ def with_options(
                 message_compression_algorithm,
                 session_options.timeouts,
                 session_options.host_health_monitor,
+                session_options.authn_credential_provider,
                 session_options.num_processing_threads,
                 session_options.blob_buffer_size,
                 session_options.channel_high_watermark,

src/blazingmq/_typing.py

@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Callable
 from typing import Mapping
+from typing import Optional
 from typing import Union
 
 from ._enums import PropertyType
@@ -23,3 +25,9 @@
 PropertyValueDict = Mapping[str, PropertyValueType]
 
 PropertyTypeDict = Mapping[str, PropertyType]
+
+AuthnCredentialProvider = Callable[[], Optional[tuple[str, bytes]]]
+"""A callable that returns authentication credentials as a tuple of
+``(mechanism, data)`` tuple of ``(str, bytes)``, or ``None`` if an
+error occurs while obtaining credentials.
+"""

src/cpp/pybmq_session.cpp

@@ -15,6 +15,7 @@
 
 #include <pybmq_session.h>
 
+#include <pybmq_gilacquireguard.h>
 #include <pybmq_gilreleaseguard.h>
 #include <pybmq_messageutils.h>
 #include <pybmq_mocksession.h>
@@ -77,13 +78,14 @@ Session::Session(
         PyObject* py_session_event_callback,
         PyObject* py_message_event_callback,
         PyObject* py_ack_event_callback,
+        PyObject* authn_credential_cb,
         const char* broker_uri,
         const char* script_name,
         bmqt::CompressionAlgorithmType::Enum message_compression_type,
         bsl::optional<int> num_processing_threads,
         bsl::optional<int> blob_buffer_size,
         bsl::optional<int> channel_high_watermark,
-        bsl::optional<bsl::pair<int, int> > event_queue_watermarks,
+        bsl::optional<bsl::pair<int, int>> event_queue_watermarks,
         const bsls::TimeInterval& stats_dump_interval,
         const bsls::TimeInterval& connect_timeout,
         const bsls::TimeInterval& disconnect_timeout,
@@ -119,6 +121,73 @@ Session::Session(
     }
 
     d_message_compression_type = message_compression_type;
+
+    bmqt::SessionOptions::AuthnCredentialCb cpp_callback;
+    bool has_auth_callback = false;
+
+    if (authn_credential_cb != nullptr && authn_credential_cb != Py_None) {
+        // Increment reference count since we're storing the Python object
+        Py_INCREF(authn_credential_cb);
+        has_auth_callback = true;
+
+        // Create a C++ lambda that wraps the Python callback
+        // TODO this can't be a lambda
+        cpp_callback =
+                [authn_credential_cb](
+                        bsl::ostream& error) -> bsl::optional<bmqt::AuthnCredential> {
+            pybmq::GilAcquireGuard guard;
+
+            // Call get_credential_data() method on the Python object
+            bslma::ManagedPtr<PyObject> result =
+                    RefUtils::toManagedPtr(PyObject_CallMethod(
+                            authn_credential_cb,
+                            "get_credential_data",
+                            nullptr));
+
+            if (!result) {
+                // Python exception occurred
+                PyErr_Print();
+                error << "Error calling get_credential_data()";
+                return bsl::optional<bmqt::AuthnCredential>();
+            }
+
+            if (result.get() == Py_None) {
+                return bsl::optional<bmqt::AuthnCredential>();
+            }
+
+            // Extract tuple (mechanism, data)
+            if (!PyTuple_Check(result.get()) || PyTuple_Size(result.get()) != 2) {
+                error << "get_credential_data() must return (str, bytes) or None";
+                return bsl::optional<bmqt::AuthnCredential>();
+            }
+
+            PyObject* mechanism_obj = PyTuple_GetItem(result.get(), 0);
+            PyObject* data_obj = PyTuple_GetItem(result.get(), 1);
+
+            if (!PyUnicode_Check(mechanism_obj) || !PyBytes_Check(data_obj)) {
+                error << "get_credential_data() must return (str, bytes) or None";
+                return bsl::optional<bmqt::AuthnCredential>();
+            }
+
+            // Convert Python str to C++ string
+            const char* mechanism_cstr = PyUnicode_AsUTF8(mechanism_obj);
+            bsl::string mechanism(mechanism_cstr);
+
+            // Convert Python bytes to vector<char>
+            char* data_ptr;
+            Py_ssize_t data_len;
+            PyBytes_AsStringAndSize(data_obj, &data_ptr, &data_len);
+            bsl::vector<char> data(data_ptr, data_ptr + data_len);
+
+            // Construct and move credential into optional
+            // (AuthnCredential is move-only)
+            bmqt::AuthnCredential credential(mechanism, data);
+            bsl::optional<bmqt::AuthnCredential> opt_credential(
+                    bslmf::MovableRefUtil::move(credential));
+            return opt_credential;
+        };
+    }
+
     {
         pybmq::GilReleaseGuard guard;
         bmqt::SessionOptions options;
@@ -144,6 +213,10 @@ Session::Session(
                     event_queue_watermarks.value().second);
         }
 
+        if (has_auth_callback) {
+            options.setAuthnCredentialCb(cpp_callback);
+        }
+
         if (stats_dump_interval != bsls::TimeInterval()) {
             options.setStatsDumpInterval(stats_dump_interval);
         }
@@ -529,8 +602,8 @@ Session::post(
             oss << "Failed to post message to " << queue_uri << " queue: " << post_rc;
             throw GenericError(oss.str());
         }
-        // We have a successful post and the SDK now owns the `on_ack` callback object
-        // so release our reference without a DECREF.
+        // We have a successful post and the SDK now owns the `on_ack` callback
+        // object so release our reference without a DECREF.
         managed_on_ack.release();
     } catch (const GenericError& exc) {
         PyErr_SetString(d_error, exc.what());