Skip to content

Commit

Permalink
Python : Client API for retrieving internal statistics (valkey-io#2707)
Browse files Browse the repository at this point in the history
* add get statistics support to python wrapper

Signed-off-by: BoazBD <boazbd@amazon.com>

* fmt rust

Signed-off-by: BoazBD <boazbd@amazon.com>

* fix deprecated funcs

Signed-off-by: BoazBD <boazbd@amazon.com>

* fmt again

Signed-off-by: BoazBD <boazbd@amazon.com>

* add to changelog

Signed-off-by: BoazBD <boazbd@amazon.com>

* refactor imports

Signed-off-by: BoazBD <boazbd@amazon.com>

---------

Signed-off-by: BoazBD <boazbd@amazon.com>
  • Loading branch information
BoazBD authored Nov 19, 2024
1 parent b8dc6a5 commit 86fa0d4
Show file tree
Hide file tree
Showing 5 changed files with 45 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#### Changes
* Python: Client API for retrieving internal statistics ([#2707](https://github.com/valkey-io/valkey-glide/pull/2707))
* Node, Python: Adding support for replacing connection configured password ([#2651](https://github.com/valkey-io/valkey-glide/pull/2651))
* Node: Add FT._ALIASLIST command([#2652](https://github.com/valkey-io/valkey-glide/pull/2652))
* Python: Python: `FT._ALIASLIST` command added([#2638](https://github.com/valkey-io/valkey-glide/pull/2638))
Expand Down
1 change: 1 addition & 0 deletions python/python/glide/glide.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -31,5 +31,6 @@ def start_socket_listener_external(init_callback: Callable) -> None: ...
def value_from_pointer(pointer: int) -> TResult: ...
def create_leaked_value(message: str) -> int: ...
def create_leaked_bytes_vec(args_vec: List[bytes]) -> int: ...
def get_statistics() -> dict: ...
def py_init(level: Optional[Level], file_name: Optional[str]) -> Level: ...
def py_log(log_level: Level, log_identifier: str, message: str) -> None: ...
4 changes: 4 additions & 0 deletions python/python/glide/glide_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
MAX_REQUEST_ARGS_LEN,
ClusterScanCursor,
create_leaked_bytes_vec,
get_statistics,
start_socket_listener_external,
value_from_pointer,
)
Expand Down Expand Up @@ -533,6 +534,9 @@ async def _reader_loop(self) -> None:
else:
await self._process_response(response=response)

async def get_statistics(self) -> dict:
return get_statistics()

async def _update_connection_password(
self, password: Optional[str], re_auth: bool
) -> TResult:
Expand Down
9 changes: 9 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,15 @@ async def test_closed_client_raises_error(self, glide_client: TGlideClient):
await glide_client.set("foo", "bar")
assert "the client is closed" in str(e)

@pytest.mark.parametrize("cluster_mode", [True, False])
@pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
async def test_statistics(self, glide_client: TGlideClient):
stats = await glide_client.get_statistics()
assert isinstance(stats, dict)
assert "total_connections" in stats
assert "total_clients" in stats
assert len(stats) == 2


@pytest.mark.asyncio
class TestCommands:
Expand Down
31 changes: 30 additions & 1 deletion python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@ use glide_core::client::FINISHED_SCAN_CURSOR;
* Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
*/
use glide_core::start_socket_listener;
use glide_core::Telemetry;
use glide_core::MAX_REQUEST_ARGS_LENGTH;
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyBool, PyBytes, PyDict, PyFloat, PyList, PySet};
use pyo3::types::{PyAny, PyBool, PyBytes, PyDict, PyFloat, PyList, PySet, PyString};
use pyo3::Python;
use redis::Value;
use std::collections::HashMap;
use std::sync::Arc;

pub const DEFAULT_TIMEOUT_IN_MILLISECONDS: u32 =
Expand Down Expand Up @@ -120,12 +122,39 @@ fn glide(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(value_from_pointer, m)?)?;
m.add_function(wrap_pyfunction!(create_leaked_value, m)?)?;
m.add_function(wrap_pyfunction!(create_leaked_bytes_vec, m)?)?;
m.add_function(wrap_pyfunction!(get_statistics, m)?)?;

#[pyfunction]
fn py_log(log_level: Level, log_identifier: String, message: String) {
log(log_level, log_identifier, message);
}

#[pyfunction]
fn get_statistics(_py: Python) -> PyResult<PyObject> {
let mut stats_map = HashMap::<String, String>::new();
stats_map.insert(
"total_connections".to_string(),
Telemetry::total_connections().to_string(),
);
stats_map.insert(
"total_clients".to_string(),
Telemetry::total_clients().to_string(),
);

Python::with_gil(|py| {
let py_dict = PyDict::new_bound(py);

for (key, value) in stats_map {
py_dict.set_item(
PyString::new_bound(py, &key),
PyString::new_bound(py, &value),
)?;
}

Ok(py_dict.into_py(py))
})
}

#[pyfunction]
#[pyo3(signature = (level=None, file_name=None))]
fn py_init(level: Option<Level>, file_name: Option<&str>) -> Level {
Expand Down

0 comments on commit 86fa0d4

Please sign in to comment.