Skip to content

Commit

Permalink
Add optional client_name property to RedisConnectionInfo, which will …
Browse files Browse the repository at this point in the history
…be used with 'CLIENT SETNAME' command during connection setup.
  • Loading branch information
ikolomi committed Dec 10, 2023
1 parent 46e8a9d commit c173cfc
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 44 deletions.
14 changes: 14 additions & 0 deletions redis/src/aio/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,20 @@ where
}
}

if connection_info.client_name.is_some() {
match cmd("CLIENT")
.arg("SETNAME")
.arg(connection_info.client_name.as_ref().unwrap())
.query_async(con).await
{
Ok(Value::Okay) => {}
_ => fail!((
ErrorKind::ResponseError,
"Redis server refused to set client name"
)),
}
}

// result is ignored, as per the command's instructions.
// https://redis.io/commands/client-setinfo/
let _: RedisResult<()> = crate::connection::client_set_info_pipeline()
Expand Down
23 changes: 23 additions & 0 deletions redis/src/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,8 @@ pub struct RedisConnectionInfo {
pub password: Option<String>,
/// Use RESP 3 mode, Redis 6 or newer is required.
pub use_resp3: bool,
/// Optionally a pass a client name that should be used for connection
pub client_name: Option<String>,
}

impl FromStr for ConnectionInfo {
Expand Down Expand Up @@ -387,6 +389,7 @@ fn url_to_tcp_connection_info(url: url::Url) -> RedisResult<ConnectionInfo> {
Some(v) => v == "true",
_ => false,
},
client_name: None
},
})
}
Expand All @@ -413,6 +416,7 @@ fn url_to_unix_connection_info(url: url::Url) -> RedisResult<ConnectionInfo> {
Some(v) => v == "true",
_ => false,
},
client_name: None
},
})
}
Expand Down Expand Up @@ -979,6 +983,20 @@ fn setup_connection(
}
}

if connection_info.client_name.is_some() {
match cmd("CLIENT")
.arg("SETNAME")
.arg(connection_info.client_name.as_ref().unwrap())
.query::<Value>(&mut rv)
{
Ok(Value::Okay) => {}
_ => fail!((
ErrorKind::ResponseError,
"Redis server refused to set client name"
)),
}
}

// result is ignored, as per the command's instructions.
// https://redis.io/commands/client-setinfo/
let _: RedisResult<()> = client_set_info_pipeline().query(&mut rv);
Expand Down Expand Up @@ -1708,6 +1726,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("#@<>$".to_string()),
use_resp3: false,
client_name: None
},
},
),
Expand Down Expand Up @@ -1775,6 +1794,7 @@ mod tests {
username: None,
password: None,
use_resp3: false,
client_name: None
},
},
),
Expand All @@ -1787,6 +1807,7 @@ mod tests {
username: None,
password: None,
use_resp3: false,
client_name: None
},
},
),
Expand All @@ -1802,6 +1823,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("#@<>$".to_string()),
use_resp3: false,
client_name: None
},
},
),
Expand All @@ -1817,6 +1839,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("&?= *+".to_string()),
use_resp3: false,
client_name: None
},
},
),
Expand Down
99 changes: 55 additions & 44 deletions redis/tests/support/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -349,28 +349,7 @@ impl TestContext {
Self::with_modules(&[], true)
}

pub fn with_tls(tls_files: TlsFilePaths, mtls_enabled: bool) -> TestContext {
let redis_port = get_random_available_port();
let addr = RedisServer::get_addr(redis_port);

let server = RedisServer::new_with_addr_tls_modules_and_spawner(
addr,
None,
Some(tls_files),
mtls_enabled,
&[],
|cmd| {
cmd.spawn()
.unwrap_or_else(|err| panic!("Failed to run {cmd:?}: {err}"))
},
);

#[cfg(feature = "tls-rustls")]
let client =
build_single_client(server.connection_info(), &server.tls_paths, mtls_enabled).unwrap();
#[cfg(not(feature = "tls-rustls"))]
let client = redis::Client::open(server.connection_info()).unwrap();

fn connect_with_retries(client: &redis::Client) {
let mut con;

let millisecond = Duration::from_millis(1);
Expand All @@ -395,6 +374,31 @@ impl TestContext {
}
}
redis::cmd("FLUSHDB").execute(&mut con);
}

pub fn with_tls(tls_files: TlsFilePaths, mtls_enabled: bool) -> TestContext {
let redis_port = get_random_available_port();
let addr: ConnectionAddr = RedisServer::get_addr(redis_port);

let server = RedisServer::new_with_addr_tls_modules_and_spawner(
addr,
None,
Some(tls_files),
mtls_enabled,
&[],
|cmd| {
cmd.spawn()
.unwrap_or_else(|err| panic!("Failed to run {cmd:?}: {err}"))
},
);

#[cfg(feature = "tls-rustls")]
let client =
build_single_client(server.connection_info(), &server.tls_paths, mtls_enabled).unwrap();
#[cfg(not(feature = "tls-rustls"))]
let client = redis::Client::open(server.connection_info()).unwrap();

Self::connect_with_retries(&client);

TestContext {
server,
Expand All @@ -412,30 +416,37 @@ impl TestContext {
#[cfg(not(feature = "tls-rustls"))]
let client = redis::Client::open(server.connection_info()).unwrap();

let mut con;
Self::connect_with_retries(&client);

let millisecond = Duration::from_millis(1);
let mut retries = 0;
loop {
match client.get_connection() {
Err(err) => {
if err.is_connection_refusal() {
sleep(millisecond);
retries += 1;
if retries > 100000 {
panic!("Tried to connect too many times, last error: {err}");
}
} else {
panic!("Could not connect: {err}");
}
}
Ok(x) => {
con = x;
break;
}
}
TestContext {
server,
client,
use_resp3: use_resp3(),
}
redis::cmd("FLUSHDB").execute(&mut con);
}

pub fn with_clientname(clientname: &str) -> TestContext {
let server = RedisServer::with_modules(&[], false);

#[cfg(feature = "tls-rustls")]
let client =
build_single_client(server.connection_info(), &server.tls_paths, mtls_enabled).unwrap();
#[cfg(not(feature = "tls-rustls"))]

let con_info = redis::ConnectionInfo {
addr: server.client_addr().clone(),
redis: redis::RedisConnectionInfo {
db: Default::default(),
username: None,
password: None,
use_resp3: Default::default(),
client_name: Some(clientname.to_string())
}
};

let client = redis::Client::open(con_info).unwrap();

Self::connect_with_retries(&client);

TestContext {
server,
Expand Down
1 change: 1 addition & 0 deletions redis/tests/test_async.rs
Original file line number Diff line number Diff line change
Expand Up @@ -472,6 +472,7 @@ async fn invalid_password_issue_343() {
username: None,
password: Some("asdcasc".to_string()),
use_resp3: false,
client_name: None
},
};
let client = redis::Client::open(coninfo).unwrap();
Expand Down
30 changes: 30 additions & 0 deletions redis/tests/test_basic.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#![allow(clippy::let_unit_value)]


use redis::{
Commands, ConnectionInfo, ConnectionLike, ControlFlow, ErrorKind, ExistenceCheck, Expiry,
PubSubCommands, RedisResult, SetExpiry, SetOptions, ToRedisArgs,
Expand Down Expand Up @@ -1448,3 +1449,32 @@ fn test_blocking_sorted_set_api() {
);
}
}

#[test]
fn test_set_client_name_by_config() {
const CLIENT_NAME: &str = "TEST_CLIENT_NAME";

let ctx = TestContext::with_clientname(CLIENT_NAME);
let mut con = ctx.connection();

let client_list: redis::InfoDict = redis::cmd("CLIENT").arg("LIST").query(&mut con).unwrap();
assert!(client_list.len() == 1, "Expecting single client, got: {}", client_list.len());

let val: &redis::Value = client_list.values().next().unwrap();
match val {
redis::Value::SimpleString(s) => {
let client_attr: Vec<&str> = s.split(' ').collect();
for i in client_attr.iter() {
let this_attr: Vec<&str> = i.split('=').collect();
if this_attr[0] == "name" {
assert!(this_attr[1] == CLIENT_NAME, "Incorrect client name, expecting: {}, got {}", CLIENT_NAME, this_attr[1]);
return;
}
}
},
_ => {
assert!(false, "Unexpected Enum type returned by CLIENT LIST command (expected String)");
}
}
assert!(false, "Could not detect the expected client name {} in CLIENT LIST output", CLIENT_NAME);
}

0 comments on commit c173cfc

Please sign in to comment.