From 935f7c0bce4bae000053e191ceeb2d9b201c4da2 Mon Sep 17 00:00:00 2001 From: honglei Date: Wed, 5 Aug 2026 19:57:30 +0800 Subject: [PATCH] Add KIP-546 client quota Admin APIs Expose DescribeClientQuotas and AlterClientQuotas through the Python AdminClient, with immutable quota entity keys, input validation, type stubs, unit coverage, and a Kafka-backed roundtrip test. --- CHANGELOG.md | 1 + src/confluent_kafka/admin/__init__.py | 106 ++++ src/confluent_kafka/admin/_client_quota.py | 91 ++++ src/confluent_kafka/cimpl.pyi | 8 + src/confluent_kafka/src/Admin.c | 483 ++++++++++++++++++ tests/integration/admin/test_client_quotas.py | 46 ++ tests/test_ClientQuota.py | 80 +++ 7 files changed, 815 insertions(+) create mode 100644 src/confluent_kafka/admin/_client_quota.py create mode 100644 tests/integration/admin/test_client_quotas.py create mode 100644 tests/test_ClientQuota.py diff --git a/CHANGELOG.md b/CHANGELOG.md index faeddfd03..ba35c67d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Enhancements +- Add KIP-546 `describe_client_quotas()` and `alter_client_quotas()` Admin APIs. - Add support for saving Azure key version with DEK (#2306) - Pass context when clients make KEK calls to DEK Registry (#2308) - Minor fix for subjectPrefix parameter in subjects API (#2311) diff --git a/src/confluent_kafka/admin/__init__.py b/src/confluent_kafka/admin/__init__.py index 0507e2f1e..4b90695fe 100644 --- a/src/confluent_kafka/admin/__init__.py +++ b/src/confluent_kafka/admin/__init__.py @@ -58,6 +58,14 @@ ) from ._acl import AclOperation # noqa: F401 from ._acl import AclBinding, AclBindingFilter, AclPermissionType # noqa: F401 +from ._client_quota import ( # noqa: F401 + ClientQuotaAlteration, + ClientQuotaAlterationOp, + ClientQuotaEntity, + ClientQuotaFilter, + ClientQuotaFilterComponent, + ClientQuotaMatchType, +) from ._cluster import DescribeClusterResult # noqa: F401 # Unused imports are keeped to be accessible using this public module @@ -528,6 +536,71 @@ def _check_alter_user_scram_credentials_request(alterations: List[UserScramCrede + "UserScramCredentialDeletion" ) + @staticmethod + def _check_client_quota_filter(client_quota_filter: ClientQuotaFilter) -> None: + if not isinstance(client_quota_filter, ClientQuotaFilter): + raise TypeError("client_quota_filter must be a ClientQuotaFilter") + if not isinstance(client_quota_filter.components, list): + raise TypeError("components must be a list") + if not isinstance(client_quota_filter.strict, bool): + raise TypeError("strict must be a bool") + + entity_types = set() + for component in client_quota_filter.components: + if not isinstance(component, ClientQuotaFilterComponent): + raise TypeError("components must contain ClientQuotaFilterComponent objects") + if not isinstance(component.entity_type, string_type): + raise TypeError("entity_type must be a string") + if not component.entity_type: + raise ValueError("entity_type must not be empty") + if component.entity_type in entity_types: + raise ValueError("entity_type values must be unique") + entity_types.add(component.entity_type) + if not isinstance(component.match_type, ClientQuotaMatchType): + raise TypeError("match_type must be a ClientQuotaMatchType") + if component.match_type is ClientQuotaMatchType.EXACT: + if not isinstance(component.match, string_type) or not component.match: + raise ValueError("EXACT match requires a non-empty match string") + elif component.match is not None: + raise ValueError("DEFAULT and ANY matches require match=None") + + @staticmethod + def _check_client_quota_alterations(alterations: List[ClientQuotaAlteration]) -> None: + if not isinstance(alterations, list): + raise TypeError("alterations must be a list") + if not alterations: + raise ValueError("alterations must not be empty") + + entities = set() + for alteration in alterations: + if not isinstance(alteration, ClientQuotaAlteration): + raise TypeError("alterations must contain ClientQuotaAlteration objects") + if not isinstance(alteration.entity, ClientQuotaEntity): + raise TypeError("entity must be a ClientQuotaEntity") + if not alteration.entity.entries: + raise ValueError("entity entries must not be empty") + for entity_type, entity_name in alteration.entity.entries.items(): + if not isinstance(entity_type, string_type) or not entity_type: + raise ValueError("entity type must be a non-empty string") + if entity_name is not None and (not isinstance(entity_name, string_type) or not entity_name): + raise ValueError("entity name must be None or a non-empty string") + if alteration.entity in entities: + raise ValueError("quota entities must be unique") + entities.add(alteration.entity) + if not isinstance(alteration.ops, list) or not alteration.ops: + raise ValueError("ops must be a non-empty list") + keys = set() + for op in alteration.ops: + if not isinstance(op, ClientQuotaAlterationOp): + raise TypeError("ops must contain ClientQuotaAlterationOp objects") + if not isinstance(op.key, string_type) or not op.key: + raise ValueError("quota key must be a non-empty string") + if op.key in keys: + raise ValueError("quota keys must be unique within an alteration") + keys.add(op.key) + if op.value is not None and not isinstance(op.value, (int, float)): + raise TypeError("quota value must be numeric or None") + @staticmethod def _check_list_offsets_request( topic_partition_offsets: Dict[_TopicPartition, OffsetSpec], kwargs: Dict[str, Any] @@ -1261,6 +1334,39 @@ def alter_user_scram_credentials( # type: ignore[override] super(AdminClient, self).alter_user_scram_credentials(alterations, f, **kwargs) return futmap + def describe_client_quotas( # type: ignore[override] + self, client_quota_filter: ClientQuotaFilter, **kwargs: Any + ) -> concurrent.futures.Future: + """Describe client quotas matching ``client_quota_filter``. + + :param ClientQuotaFilter client_quota_filter: Quota entity filter. + :param float request_timeout: Overall request timeout in seconds. + :returns: A future yielding ``dict[ClientQuotaEntity, dict[str, float]]``. + """ + AdminClient._check_client_quota_filter(client_quota_filter) + internal_f, result_f = AdminClient._make_single_future_pair() + super(AdminClient, self).describe_client_quotas(client_quota_filter, internal_f, **kwargs) + return result_f + + def alter_client_quotas( # type: ignore[override] + self, alterations: List[ClientQuotaAlteration], **kwargs: Any + ) -> Dict[ClientQuotaEntity, concurrent.futures.Future]: + """Alter client quotas for one or more entities. + + :param list(ClientQuotaAlteration) alterations: Quota alterations. + :param bool validate_only: Validate without applying changes. + :param float request_timeout: Overall request timeout in seconds. + :returns: Futures keyed by client quota entity. + """ + AdminClient._check_client_quota_alterations(alterations) + f, futmap = AdminClient._make_futures_v2( + [alteration.entity for alteration in alterations], + ClientQuotaEntity, + AdminClient._make_futmap_result, + ) + super(AdminClient, self).alter_client_quotas(alterations, f, **kwargs) + return futmap + def list_offsets( # type: ignore[override] self, topic_partition_offsets: Dict[_TopicPartition, OffsetSpec], **kwargs: Any ) -> Dict[_TopicPartition, concurrent.futures.Future]: diff --git a/src/confluent_kafka/admin/_client_quota.py b/src/confluent_kafka/admin/_client_quota.py new file mode 100644 index 000000000..ad9b63c4c --- /dev/null +++ b/src/confluent_kafka/admin/_client_quota.py @@ -0,0 +1,91 @@ +# Copyright 2026 Confluent Inc. +# +# Licensed 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. + +from enum import Enum +from typing import Dict, List, Optional + + +class ClientQuotaMatchType(Enum): + """Match type for a client quota filter component.""" + + EXACT = 0 + DEFAULT = 1 + ANY = 2 + + +class ClientQuotaEntity: + """A client quota entity, represented by entity type and name pairs. + + A ``None`` name identifies the default entity for that entity type. + The returned ``entries`` dictionary is a copy so the entity remains a + stable dictionary key. + """ + + def __init__(self, entries: Dict[str, Optional[str]]) -> None: + self._entries = dict(entries) + + @property + def entries(self) -> Dict[str, Optional[str]]: + return dict(self._entries) + + def __eq__(self, other: object) -> bool: + return isinstance(other, ClientQuotaEntity) and self._entries == other._entries + + def __hash__(self) -> int: + return hash(frozenset(self._entries.items())) + + def __repr__(self) -> str: + return "ClientQuotaEntity({!r})".format(self._entries) + + +class ClientQuotaFilterComponent: + """One entity component in a client quota filter.""" + + def __init__( + self, + entity_type: str, + match_type: ClientQuotaMatchType, + match: Optional[str] = None, + ) -> None: + self.entity_type = entity_type + self.match_type = match_type + self.match = match + + +class ClientQuotaFilter: + """Filter used by :meth:`AdminClient.describe_client_quotas`. + + When ``strict`` is true, matching entities cannot contain entity types + absent from ``components``. + """ + + def __init__(self, components: List[ClientQuotaFilterComponent], strict: bool = False) -> None: + self.components = components + self.strict = strict + + +class ClientQuotaAlterationOp: + """A quota key alteration. ``value=None`` removes the quota key.""" + + def __init__(self, key: str, value: Optional[float]) -> None: + self.key = key + self.value = value + + +class ClientQuotaAlteration: + """A set of quota operations for one client quota entity.""" + + def __init__(self, entity: ClientQuotaEntity, ops: List[ClientQuotaAlterationOp]) -> None: + self.entity = entity + self.ops = ops diff --git a/src/confluent_kafka/cimpl.pyi b/src/confluent_kafka/cimpl.pyi index 7dd339748..b299013c4 100644 --- a/src/confluent_kafka/cimpl.pyi +++ b/src/confluent_kafka/cimpl.pyi @@ -699,6 +699,14 @@ class _AdminClientImpl: def alter_user_scram_credentials( self, alterations: List[Any], future: Any, request_timeout: float = -1 # List[UserScramCredentialAlteration] ) -> None: ... + def describe_client_quotas(self, client_quota_filter: Any, future: Any, request_timeout: float = -1) -> None: ... + def alter_client_quotas( + self, + alterations: List[Any], + future: Any, + validate_only: bool = False, + request_timeout: float = -1, + ) -> None: ... def list_offsets( self, topic_partitions: List[TopicPartition], diff --git a/src/confluent_kafka/src/Admin.c b/src/confluent_kafka/src/Admin.c index 0c8a97070..a7d8fd5ab 100644 --- a/src/confluent_kafka/src/Admin.c +++ b/src/confluent_kafka/src/Admin.c @@ -3409,6 +3409,318 @@ const char Admin_elect_leaders_doc[] = PyDoc_STR( " This method should not be used directly, use " "confluent_kafka.AdminClient.elect_leaders()\n"); + +static PyObject * +Admin_describe_client_quotas(Handle *self, PyObject *args, PyObject *kwargs) { + PyObject *filter, *future, *components = NULL, *strict_obj = NULL; + static char *kws[] = {"client_quota_filter", "future", + "request_timeout", NULL}; + struct Admin_options options = Admin_options_INITIALIZER; + rd_kafka_AdminOptions_t *c_options = NULL; + rd_kafka_ClientQuotaFilter_t *c_filter = NULL; + rd_kafka_queue_t *rkqu = NULL; + Py_ssize_t i, component_cnt; + int strict; + char errstr[512]; + CallState cs; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|f", kws, &filter, + &future, &options.request_timeout)) + return NULL; + + components = PyObject_GetAttrString(filter, "components"); + strict_obj = PyObject_GetAttrString(filter, "strict"); + if (!components || !strict_obj || !PyList_Check(components)) { + PyErr_SetString(PyExc_TypeError, + "client_quota_filter must provide list " + "components and bool strict attributes"); + goto err; + } + + strict = PyObject_IsTrue(strict_obj); + if (strict == -1) + goto err; + c_filter = rd_kafka_ClientQuotaFilter_new(strict); + if (!c_filter) { + PyErr_NoMemory(); + goto err; + } + component_cnt = PyList_Size(components); + for (i = 0; i < component_cnt; i++) { + PyObject *component = PyList_GET_ITEM(components, i); + PyObject *type_obj = NULL, *match_type_obj = NULL; + PyObject *match_value_obj = NULL, *match_obj = NULL; + const char *entity_type, *match = NULL; + long match_type; + rd_kafka_resp_err_t c_err; + + type_obj = PyObject_GetAttrString(component, "entity_type"); + match_type_obj = + PyObject_GetAttrString(component, "match_type"); + match_obj = PyObject_GetAttrString(component, "match"); + if (match_type_obj) + match_value_obj = + PyObject_GetAttrString(match_type_obj, "value"); + if (!type_obj || !match_type_obj || !match_value_obj || + !match_obj || !PyUnicode_Check(type_obj) || + (match_obj != Py_None && !PyUnicode_Check(match_obj))) { + Py_XDECREF(type_obj); + Py_XDECREF(match_type_obj); + Py_XDECREF(match_value_obj); + Py_XDECREF(match_obj); + PyErr_Format(PyExc_TypeError, + "Invalid client quota filter component " + "at index %zd", + i); + goto err; + } + + entity_type = PyUnicode_AsUTF8(type_obj); + match_type = PyLong_AsLong(match_value_obj); + if (match_obj != Py_None) + match = PyUnicode_AsUTF8(match_obj); + if (!entity_type || (match_obj != Py_None && !match) || + PyErr_Occurred()) { + Py_DECREF(type_obj); + Py_DECREF(match_type_obj); + Py_DECREF(match_value_obj); + Py_DECREF(match_obj); + goto err; + } + c_err = rd_kafka_ClientQuotaFilter_add_component( + c_filter, entity_type, + (rd_kafka_ClientQuotaMatchType_t)match_type, match, errstr, + sizeof(errstr)); + Py_DECREF(type_obj); + Py_DECREF(match_type_obj); + Py_DECREF(match_value_obj); + Py_DECREF(match_obj); + if (c_err) { + PyErr_SetString(PyExc_ValueError, errstr); + goto err; + } + } + + c_options = Admin_options_to_c( + self, RD_KAFKA_ADMIN_OP_DESCRIBECLIENTQUOTAS, &options, future); + if (!c_options) + goto err; + Py_INCREF(future); + rkqu = rd_kafka_queue_get_background(self->rk); + CallState_begin(self, &cs); + rd_kafka_DescribeClientQuotas(self->rk, c_filter, c_options, rkqu); + CallState_end(self, &cs); + + rd_kafka_queue_destroy(rkqu); + rd_kafka_AdminOptions_destroy(c_options); + rd_kafka_ClientQuotaFilter_destroy(c_filter); + Py_DECREF(components); + Py_DECREF(strict_obj); + Py_RETURN_NONE; + +err: + if (rkqu) + rd_kafka_queue_destroy(rkqu); + if (c_options) { + rd_kafka_AdminOptions_destroy(c_options); + Py_DECREF(future); + } + if (c_filter) + rd_kafka_ClientQuotaFilter_destroy(c_filter); + Py_XDECREF(components); + Py_XDECREF(strict_obj); + return NULL; +} + + +static PyObject * +Admin_alter_client_quotas(Handle *self, PyObject *args, PyObject *kwargs) { + PyObject *alterations, *future; + static char *kws[] = {"alterations", "future", "validate_only", + "request_timeout", NULL}; + struct Admin_options options = Admin_options_INITIALIZER; + rd_kafka_AdminOptions_t *c_options = NULL; + rd_kafka_ClientQuotaEntry_t **c_entries = NULL; + rd_kafka_queue_t *rkqu = NULL; + Py_ssize_t i, alteration_cnt; + size_t c_entry_cnt = 0; + char errstr[512]; + CallState cs; + + if (!PyArg_ParseTupleAndKeywords( + args, kwargs, "OO|pf", kws, &alterations, &future, + &options.validate_only, &options.request_timeout)) + return NULL; + if (!PyList_Check(alterations) || PyList_Size(alterations) == 0) { + PyErr_SetString(PyExc_ValueError, + "alterations must be a non-empty list"); + return NULL; + } + + alteration_cnt = PyList_Size(alterations); + c_entries = calloc((size_t)alteration_cnt, sizeof(*c_entries)); + if (!c_entries) + return PyErr_NoMemory(); + for (i = 0; i < alteration_cnt; i++) { + PyObject *alteration = PyList_GET_ITEM(alterations, i); + PyObject *entity = NULL, *entries = NULL, *ops = NULL; + Py_ssize_t pos = 0, j, op_cnt; + PyObject *entity_type_obj, *entity_name_obj; + + c_entries[i] = rd_kafka_ClientQuotaEntry_new(); + if (!c_entries[i]) { + PyErr_NoMemory(); + goto err; + } + c_entry_cnt++; + entity = PyObject_GetAttrString(alteration, "entity"); + if (entity) + entries = PyObject_GetAttrString(entity, "entries"); + ops = PyObject_GetAttrString(alteration, "ops"); + if (!entity || !entries || !PyDict_Check(entries) || !ops || + !PyList_Check(ops)) { + Py_XDECREF(entity); + Py_XDECREF(entries); + Py_XDECREF(ops); + PyErr_Format(PyExc_TypeError, + "Invalid client quota alteration at " + "index %zd", + i); + goto err; + } + + while (PyDict_Next(entries, &pos, &entity_type_obj, + &entity_name_obj)) { + const char *entity_type, *entity_name = NULL; + rd_kafka_resp_err_t c_err; + + if (!PyUnicode_Check(entity_type_obj) || + (entity_name_obj != Py_None && + !PyUnicode_Check(entity_name_obj))) { + PyErr_SetString( + PyExc_TypeError, + "quota entity entries must map strings to " + "strings or None"); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + entity_type = PyUnicode_AsUTF8(entity_type_obj); + if (entity_name_obj != Py_None) + entity_name = PyUnicode_AsUTF8(entity_name_obj); + if (!entity_type || + (entity_name_obj != Py_None && !entity_name)) { + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + c_err = rd_kafka_ClientQuotaEntry_add_entity( + c_entries[i], entity_type, entity_name, errstr, + sizeof(errstr)); + if (c_err) { + PyErr_SetString(PyExc_ValueError, errstr); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + } + + op_cnt = PyList_Size(ops); + for (j = 0; j < op_cnt; j++) { + PyObject *op = PyList_GET_ITEM(ops, j); + PyObject *key_obj = PyObject_GetAttrString(op, "key"); + PyObject *value_obj = + PyObject_GetAttrString(op, "value"); + const char *key; + double value = 0.0; + int remove; + rd_kafka_resp_err_t c_err; + + if (!key_obj || !value_obj || + !PyUnicode_Check(key_obj)) { + Py_XDECREF(key_obj); + Py_XDECREF(value_obj); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + PyErr_SetString(PyExc_TypeError, + "Invalid client quota op"); + goto err; + } + key = PyUnicode_AsUTF8(key_obj); + if (!key) { + Py_DECREF(key_obj); + Py_DECREF(value_obj); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + remove = value_obj == Py_None; + if (!remove) { + value = PyFloat_AsDouble(value_obj); + if (PyErr_Occurred()) { + Py_DECREF(key_obj); + Py_DECREF(value_obj); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + } + c_err = rd_kafka_ClientQuotaEntry_add_operation( + c_entries[i], key, value, remove, errstr, + sizeof(errstr)); + Py_DECREF(key_obj); + Py_DECREF(value_obj); + if (c_err) { + PyErr_SetString(PyExc_ValueError, errstr); + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + goto err; + } + } + Py_DECREF(entity); + Py_DECREF(entries); + Py_DECREF(ops); + } + + c_options = Admin_options_to_c( + self, RD_KAFKA_ADMIN_OP_ALTERCLIENTQUOTAS, &options, future); + if (!c_options) + goto err; + Py_INCREF(future); + rkqu = rd_kafka_queue_get_background(self->rk); + CallState_begin(self, &cs); + rd_kafka_AlterClientQuotas(self->rk, c_entries, (size_t)alteration_cnt, + c_options, rkqu); + CallState_end(self, &cs); + + rd_kafka_queue_destroy(rkqu); + rd_kafka_AdminOptions_destroy(c_options); + rd_kafka_ClientQuotaEntry_destroy_array(c_entries, + (size_t)alteration_cnt); + free(c_entries); + Py_RETURN_NONE; + +err: + if (rkqu) + rd_kafka_queue_destroy(rkqu); + if (c_options) { + rd_kafka_AdminOptions_destroy(c_options); + Py_DECREF(future); + } + if (c_entries) { + rd_kafka_ClientQuotaEntry_destroy_array(c_entries, c_entry_cnt); + free(c_entries); + } + return NULL; +} + /** * @brief Call rd_kafka_poll() and keep track of crashing callbacks. * @returns -1 if callback crashed (or poll() failed), else the number @@ -3606,6 +3918,12 @@ static PyMethodDef Admin_methods[] = { {"elect_leaders", (PyCFunction)Admin_elect_leaders, METH_VARARGS | METH_KEYWORDS, Admin_elect_leaders_doc}, + {"describe_client_quotas", (PyCFunction)Admin_describe_client_quotas, + METH_VARARGS | METH_KEYWORDS, + "Describe client quotas. Use AdminClient.describe_client_quotas()."}, + {"alter_client_quotas", (PyCFunction)Admin_alter_client_quotas, + METH_VARARGS | METH_KEYWORDS, + "Alter client quotas. Use AdminClient.alter_client_quotas()."}, {"__enter__", (PyCFunction)Admin_enter, METH_NOARGS, "Context manager entry."}, {"__exit__", (PyCFunction)Admin_exit, METH_VARARGS, @@ -5034,6 +5352,145 @@ static PyObject *Admin_c_DeletedRecords_to_py( return NULL; } + +static PyObject * +Admin_c_ClientQuotaEntity_to_py(const rd_kafka_ClientQuotaEntity_t **c_entities, + size_t c_entity_cnt) { + PyObject *entity_type = NULL, *entries = NULL, *entity = NULL; + size_t i; + + entity_type = + cfl_PyObject_lookup("confluent_kafka.admin", "ClientQuotaEntity"); + if (!entity_type) + return NULL; + entries = PyDict_New(); + if (!entries) + goto err; + + for (i = 0; i < c_entity_cnt; i++) { + const char *type = + rd_kafka_ClientQuotaEntity_type(c_entities[i]); + const char *name = + rd_kafka_ClientQuotaEntity_name(c_entities[i]); + PyObject *py_name = + name ? cfl_PyUnistr(_FromString(name)) : Py_None; + + if (!name) + Py_INCREF(Py_None); + if (!py_name || + PyDict_SetItemString(entries, type, py_name) == -1) { + Py_XDECREF(py_name); + goto err; + } + Py_DECREF(py_name); + } + + entity = PyObject_CallFunctionObjArgs(entity_type, entries, NULL); +err: + Py_XDECREF(entity_type); + Py_XDECREF(entries); + return entity; +} + + +static PyObject *Admin_c_AlterClientQuotasResult_to_py( + const rd_kafka_ClientQuotaEntry_t **c_entries, + size_t c_entry_cnt) { + PyObject *result = PyDict_New(); + size_t i; + + if (!result) + return NULL; + for (i = 0; i < c_entry_cnt; i++) { + const rd_kafka_ClientQuotaEntity_t **c_entities; + const rd_kafka_error_t *c_error; + size_t c_entity_cnt; + PyObject *entity, *error; + + c_entities = rd_kafka_ClientQuotaEntry_entities(c_entries[i], + &c_entity_cnt); + entity = + Admin_c_ClientQuotaEntity_to_py(c_entities, c_entity_cnt); + if (!entity) + goto err; + c_error = rd_kafka_ClientQuotaEntry_error(c_entries[i]); + error = + c_error + ? KafkaError_new_or_None(rd_kafka_error_code(c_error), + rd_kafka_error_string(c_error)) + : KafkaError_new_or_None(RD_KAFKA_RESP_ERR_NO_ERROR, + NULL); + if (!error || PyDict_SetItem(result, entity, error) == -1) { + Py_DECREF(entity); + Py_XDECREF(error); + goto err; + } + Py_DECREF(entity); + Py_DECREF(error); + } + return result; +err: + Py_DECREF(result); + return NULL; +} + + +static PyObject *Admin_c_DescribeClientQuotasResult_to_py( + const rd_kafka_DescribeClientQuotas_result_entry_t **c_entries, + size_t c_entry_cnt) { + PyObject *result = PyDict_New(); + size_t i; + + if (!result) + return NULL; + for (i = 0; i < c_entry_cnt; i++) { + const rd_kafka_ClientQuotaEntity_t **c_entities; + const rd_kafka_ClientQuotaValue_t **c_values; + size_t c_entity_cnt, c_value_cnt, j; + PyObject *entity, *values; + + c_entities = + rd_kafka_DescribeClientQuotas_result_entry_entities( + c_entries[i], &c_entity_cnt); + c_values = rd_kafka_DescribeClientQuotas_result_entry_values( + c_entries[i], &c_value_cnt); + entity = + Admin_c_ClientQuotaEntity_to_py(c_entities, c_entity_cnt); + values = PyDict_New(); + if (!entity || !values) { + Py_XDECREF(entity); + Py_XDECREF(values); + goto err; + } + for (j = 0; j < c_value_cnt; j++) { + PyObject *value = PyFloat_FromDouble( + rd_kafka_ClientQuotaValue_value(c_values[j])); + if (!value || + PyDict_SetItemString( + values, + rd_kafka_ClientQuotaValue_key(c_values[j]), + value) == -1) { + Py_XDECREF(value); + Py_DECREF(entity); + Py_DECREF(values); + goto err; + } + Py_DECREF(value); + } + if (PyDict_SetItem(result, entity, values) == -1) { + Py_DECREF(entity); + Py_DECREF(values); + goto err; + } + Py_DECREF(entity); + Py_DECREF(values); + } + return result; +err: + Py_DECREF(result); + return NULL; +} + /** * @brief Event callback triggered from librdkafka's background thread * when Admin API results are ready. @@ -5411,6 +5868,32 @@ static void Admin_background_event_cb(rd_kafka_t *rk, break; } + case RD_KAFKA_EVENT_ALTERCLIENTQUOTAS_RESULT: { + const rd_kafka_AlterClientQuotas_result_t *c_result; + const rd_kafka_ClientQuotaEntry_t **c_entries; + size_t c_entry_cnt; + + c_result = rd_kafka_event_AlterClientQuotas_result(rkev); + c_entries = rd_kafka_AlterClientQuotas_result_entries( + c_result, &c_entry_cnt); + result = Admin_c_AlterClientQuotasResult_to_py(c_entries, + c_entry_cnt); + break; + } + + case RD_KAFKA_EVENT_DESCRIBECLIENTQUOTAS_RESULT: { + const rd_kafka_DescribeClientQuotas_result_t *c_result; + const rd_kafka_DescribeClientQuotas_result_entry_t **c_entries; + size_t c_entry_cnt; + + c_result = rd_kafka_event_DescribeClientQuotas_result(rkev); + c_entries = rd_kafka_DescribeClientQuotas_result_entries( + c_result, &c_entry_cnt); + result = Admin_c_DescribeClientQuotasResult_to_py(c_entries, + c_entry_cnt); + break; + } + default: Py_DECREF(error); /* Py_None */ error = KafkaError_new0(RD_KAFKA_RESP_ERR__UNSUPPORTED_FEATURE, diff --git a/tests/integration/admin/test_client_quotas.py b/tests/integration/admin/test_client_quotas.py new file mode 100644 index 000000000..84ae43bf4 --- /dev/null +++ b/tests/integration/admin/test_client_quotas.py @@ -0,0 +1,46 @@ +import time +import uuid + +from confluent_kafka.admin import ( + ClientQuotaAlteration, + ClientQuotaAlterationOp, + ClientQuotaEntity, + ClientQuotaFilter, + ClientQuotaFilterComponent, + ClientQuotaMatchType, +) + + +def _describe_user(admin_client, user): + quota_filter = ClientQuotaFilter([ClientQuotaFilterComponent("user", ClientQuotaMatchType.EXACT, user)]) + return admin_client.describe_client_quotas(quota_filter).result(timeout=15) + + +def test_client_quotas(kafka_cluster): + admin_client = kafka_cluster.admin() + user = "confluent-kafka-python-{}".format(uuid.uuid4().hex) + entity = ClientQuotaEntity({"user": user}) + key = "producer_byte_rate" + + validate_only = ClientQuotaAlteration(entity, [ClientQuotaAlterationOp(key, 111111.0)]) + assert admin_client.alter_client_quotas([validate_only], validate_only=True)[entity].result(timeout=15) is None + assert entity not in _describe_user(admin_client, user) + + alteration = ClientQuotaAlteration(entity, [ClientQuotaAlterationOp(key, 222222.0)]) + assert admin_client.alter_client_quotas([alteration])[entity].result(timeout=15) is None + + for _ in range(6): + result = _describe_user(admin_client, user) + if result.get(entity, {}).get(key) == 222222.0: + break + time.sleep(0.2) + assert result[entity][key] == 222222.0 + + removal = ClientQuotaAlteration(entity, [ClientQuotaAlterationOp(key, None)]) + assert admin_client.alter_client_quotas([removal])[entity].result(timeout=15) is None + for _ in range(6): + result = _describe_user(admin_client, user) + if entity not in result: + break + time.sleep(0.2) + assert entity not in result diff --git a/tests/test_ClientQuota.py b/tests/test_ClientQuota.py new file mode 100644 index 000000000..99fe2b3f7 --- /dev/null +++ b/tests/test_ClientQuota.py @@ -0,0 +1,80 @@ +import pytest + +from confluent_kafka.admin import ( + AdminClient, + ClientQuotaAlteration, + ClientQuotaAlterationOp, + ClientQuotaEntity, + ClientQuotaFilter, + ClientQuotaFilterComponent, + ClientQuotaMatchType, +) + + +def test_client_quota_entity_equality_and_hash(): + first = ClientQuotaEntity({"user": "alice", "client-id": None}) + second = ClientQuotaEntity({"client-id": None, "user": "alice"}) + assert first == second + assert hash(first) == hash(second) + + +def test_client_quota_entity_entries_cannot_mutate_hash_key(): + entity = ClientQuotaEntity({"user": "alice"}) + keyed = {entity: "result"} + + entries = entity.entries + entries["user"] = "bob" + + assert entity.entries == {"user": "alice"} + assert keyed[entity] == "result" + + +def test_client_quota_filter_validation(): + valid = ClientQuotaFilter([ClientQuotaFilterComponent("user", ClientQuotaMatchType.EXACT, "alice")], strict=True) + AdminClient._check_client_quota_filter(valid) + + with pytest.raises(ValueError): + AdminClient._check_client_quota_filter( + ClientQuotaFilter([ClientQuotaFilterComponent("user", ClientQuotaMatchType.EXACT)]) + ) + with pytest.raises(ValueError): + AdminClient._check_client_quota_filter( + ClientQuotaFilter([ClientQuotaFilterComponent("user", ClientQuotaMatchType.ANY, "alice")]) + ) + with pytest.raises(ValueError): + AdminClient._check_client_quota_filter( + ClientQuotaFilter( + [ + ClientQuotaFilterComponent("user", ClientQuotaMatchType.ANY), + ClientQuotaFilterComponent("user", ClientQuotaMatchType.DEFAULT), + ] + ) + ) + + +def test_client_quota_alteration_validation(): + entity = ClientQuotaEntity({"user": "alice"}) + valid = [ClientQuotaAlteration(entity, [ClientQuotaAlterationOp("producer_byte_rate", 1024.0)])] + AdminClient._check_client_quota_alterations(valid) + + with pytest.raises(ValueError): + AdminClient._check_client_quota_alterations([]) + with pytest.raises(ValueError): + AdminClient._check_client_quota_alterations( + [ + ClientQuotaAlteration(entity, [ClientQuotaAlterationOp("producer_byte_rate", 1.0)]), + ClientQuotaAlteration(entity, [ClientQuotaAlterationOp("consumer_byte_rate", 1.0)]), + ] + ) + with pytest.raises(ValueError): + AdminClient._check_client_quota_alterations( + [ + ClientQuotaAlteration( + entity, + [ + ClientQuotaAlterationOp("producer_byte_rate", 1.0), + ClientQuotaAlterationOp("producer_byte_rate", None), + ], + ) + ] + )