diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index a6a6c124..e271e30c 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -25,6 +25,7 @@ "description": "General settings are relevant to the operations of the pooler itself, or apply to all database pools.\n\n", "$ref": "#/$defs/General", "default": { + "application_name_add_host": false, "auth_type": "scram", "ban_replica_lag": 9223372036854775807, "ban_replica_lag_bytes": 9223372036854775807, @@ -653,6 +654,11 @@ "description": "General settings are relevant to the operations of the pooler itself, or apply to all database pools.\n\n", "type": "object", "properties": { + "application_name_add_host": { + "description": "Add the client host address and port to `application_name` at connection start\nand whenever the client later changes it with `SET` / `set_config`.\n\nThe result is `{application_name} - {ip}:{port}`. If the client sent no name,\nthe prefix is empty (` - 10.0.0.5:1234`). A later `SET application_name` replaces\nthe name and the host is appended again.\n\n_Default:_ `false`\n\n", + "type": "boolean", + "default": false + }, "auth_type": { "description": "What kind of authentication mechanism to use for client connections.\n\n_Default:_ `scram`\n\n", "$ref": "#/$defs/AuthType", diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index 90d0f4e3..3ec70522 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -606,6 +606,19 @@ pub struct General { #[serde(default = "General::log_disconnections")] pub log_disconnections: bool, + /// Add the client host address and port to `application_name` at connection start + /// and whenever the client later changes it with `SET` / `set_config`. + /// + /// The result is `{application_name} - {ip}:{port}`. If the client sent no name, + /// the prefix is empty (` - 10.0.0.5:1234`). A later `SET application_name` replaces + /// the name and the host is appended again. + /// + /// _Default:_ `false` + /// + /// + #[serde(default = "General::application_name_add_host")] + pub application_name_add_host: bool, + /// Window, in milliseconds, over which to deduplicate identical log messages. Set to `0` to disable throttling. /// /// **Note:** When enabled, identical messages (same level, target, and body) that exceed `log_dedup_threshold` within this window are suppressed and replaced with a single summary line at the end of the window. @@ -955,6 +968,7 @@ impl Default for General { log_level: Self::log_level(), log_connections: Self::log_connections(), log_disconnections: Self::log_disconnections(), + application_name_add_host: Self::application_name_add_host(), log_dedup_window: 0, log_dedup_threshold: 0, two_phase_commit: bool::default(), @@ -1488,6 +1502,10 @@ impl General { Self::env_bool_or_default("PGDOG_LOG_DISCONNECTIONS", true) } + pub fn application_name_add_host() -> bool { + Self::env_bool_or_default("PGDOG_APPLICATION_NAME_ADD_HOST", false) + } + pub fn expanded_explain() -> bool { Self::env_bool_or_default("PGDOG_EXPANDED_EXPLAIN", false) } @@ -2052,21 +2070,25 @@ mod tests { let _guard = set_env_var("PGDOG_CROSS_SHARD_DISABLED", "yes"); let _guard = set_env_var("PGDOG_LOG_CONNECTIONS", "false"); let _guard = set_env_var("PGDOG_LOG_DISCONNECTIONS", "0"); + let _guard = set_env_var("PGDOG_APPLICATION_NAME_ADD_HOST", "true"); assert!(General::dry_run()); assert!(General::cross_shard_disabled()); assert!(!General::log_connections()); assert!(!General::log_disconnections()); + assert!(General::application_name_add_host()); let _guard = remove_env_var("PGDOG_DRY_RUN"); let _guard = remove_env_var("PGDOG_CROSS_SHARD_DISABLED"); let _guard = remove_env_var("PGDOG_LOG_CONNECTIONS"); let _guard = remove_env_var("PGDOG_LOG_DISCONNECTIONS"); + let _guard = remove_env_var("PGDOG_APPLICATION_NAME_ADD_HOST"); assert!(!General::dry_run()); assert!(!General::cross_shard_disabled()); assert!(General::log_connections()); assert!(General::log_disconnections()); + assert!(!General::application_name_add_host()); } #[test] diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 50d6932d..e0412172 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use pgdog_config::users::PasswordKind; -use timeouts::Timeouts; use tokio::{select, spawn}; use tracing::{Level as LogLevel, debug, enabled, error, info, trace, warn}; @@ -29,16 +28,21 @@ use crate::net::messages::{ Authentication, BackendKeyData, ErrorResponse, FromBytes, FrontendPid, Message, Password, Protocol, ProtocolVersion, ReadyForQuery, ToBytes, }; -use crate::net::{MessageBuffer, ProtocolMessage, Stream, parameter::Parameters}; +use crate::net::{ + MessageBuffer, ProtocolMessage, Stream, + parameter::{Parameters, application_name_with_host}, +}; use crate::state::State; use crate::stats::memory::MemoryUsage; use crate::util::{safe_timeout, user_database_from_params}; pub mod query_engine; +pub(crate) mod request_settings; pub mod sticky; pub mod timeouts; pub mod transaction_type; +pub(crate) use request_settings::ClientRequestSettings; pub(crate) use sticky::Sticky; pub use transaction_type::TransactionType; @@ -73,10 +77,8 @@ pub struct Client { prepared_statements: PreparedStatements, // Client transaction state. transaction: Option, - // Current timeouts to use for client/server communication. - // These change based on client state, e.g. if client is running query, - // the `query_timeout` is active, and if the client is idle, the `client_idle_timeout` is. - timeouts: Timeouts, + // Per-request settings snapshot, refreshed in [`Self::buffer`]. + request_settings: ClientRequestSettings, // Stateful buffer containing the current whole client request. // This can be a query or just a `Parse` and `Flush`, but in either case, the client // will expect a response immediately and we need to handle it. @@ -89,10 +91,6 @@ pub struct Client { sticky: Sticky, /// Client database. database: String, - /// Log queries to stdout. - query_log_stdout: bool, - /// Maximum query message size before a warning is logged. - query_size_limit: Option, } /// Inputs to the per-user client certificate check. @@ -242,7 +240,7 @@ impl Client { /// Create new frontend client from the given TCP stream. async fn login( mut stream: Stream, - params: Parameters, + mut params: Parameters, addr: SocketAddr, config: Arc, protocol_version: ProtocolVersion, @@ -253,6 +251,11 @@ impl Client { return Ok(None); } + Self::maybe_add_application_name_host( + &mut params, + addr, + config.config.general.application_name_add_host, + ); let (user, database) = user_database_from_params(¶ms); let admin = database == config.config.admin.name && config.config.admin.user == user; let admin_password = &config.config.admin.password; @@ -420,7 +423,7 @@ impl Client { params: params.clone(), prepared_statements: PreparedStatements::new(), transaction: None, - timeouts: Timeouts::from_config(&config.config.general), + request_settings: ClientRequestSettings::from_general(&config.config.general), client_request: ClientRequest::default(), stream_buffer: MessageBuffer::new( config.config.memory.message_buffer, @@ -428,11 +431,21 @@ impl Client { ), sticky: Sticky::from_params(¶ms), database: database.to_string(), - query_log_stdout: false, - query_size_limit: None, })) } + fn maybe_add_application_name_host(params: &mut Parameters, addr: SocketAddr, enabled: bool) { + if !enabled { + return; + } + + let current = params.get_default("application_name", ""); + params.insert( + "application_name", + application_name_with_host(current, &addr.to_string()), + ); + } + #[cfg(test)] pub fn new_test(stream: Stream, params: Parameters) -> Self { use crate::config::config; @@ -442,6 +455,13 @@ impl Client { connect_params.insert("database", "pgdog"); connect_params.merge(params); + let addr = SocketAddr::from(([127, 0, 0, 1], 1234)); + Self::maybe_add_application_name_host( + &mut connect_params, + addr, + config().config.general.application_name_add_host, + ); + let id = FrontendPid::new(); let key = BackendKeyData::new_frontend(ProtocolVersion::V3_0, id); let mut prepared_statements = PreparedStatements::new(); @@ -449,14 +469,14 @@ impl Client { Self { stream, - addr: SocketAddr::from(([127, 0, 0, 1], 1234)), + addr, key, comms: ClientComms::new(id), streaming: false, prepared_statements, admin: false, transaction: None, - timeouts: Timeouts::from_config(&config().config.general), + request_settings: ClientRequestSettings::from_general(&config().config.general), client_request: ClientRequest::default(), stream_buffer: MessageBuffer::new( 4096, @@ -465,8 +485,6 @@ impl Client { sticky: Sticky::from_params(&connect_params), params: connect_params, database: "pgdog".to_string(), - query_log_stdout: false, - query_size_limit: None, } } @@ -633,14 +651,13 @@ impl Client { let config = config::config(); // Configure prepared statements cache. self.prepared_statements.level = config.prepared_statements(); - self.timeouts = Timeouts::from_config(&config.config.general); - self.query_log_stdout = config.config.general.query_log_stdout; - self.query_size_limit = config.config.general.query_size_limit; + self.request_settings = ClientRequestSettings::from_general(&config.config.general); self.stream_buffer - .set_size_limit_block(config.config.general.frontend_query_size_limit_block()); + .set_size_limit_block(self.request_settings.frontend_query_size_limit_block); while !self.client_request.is_complete() { let idle_timeout = self + .request_settings .timeouts .client_idle_timeout(&state, &self.client_request); @@ -736,7 +753,7 @@ impl MemoryUsage for Client { + std::mem::size_of::() + std::mem::size_of::() * 5 + self.prepared_statements.memory_used() - + std::mem::size_of::() + + std::mem::size_of::() + self.stream_buffer.capacity() + self.client_request.memory_usage() } diff --git a/pgdog/src/frontend/client/query_engine/connect.rs b/pgdog/src/frontend/client/query_engine/connect.rs index 8c31c63d..d2a0c2de 100644 --- a/pgdog/src/frontend/client/query_engine/connect.rs +++ b/pgdog/src/frontend/client/query_engine/connect.rs @@ -42,7 +42,10 @@ impl QueryEngine { self.stats.connected(); self.debug_connected(context, false); - let query_timeout = context.timeouts.query_timeout(&self.stats.state); + let query_timeout = context + .request_settings + .timeouts + .query_timeout(&self.stats.state); let begin_stmt = self.begin_stmt.take(); // We may need to sync params with the server and that reads from the socket. diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs index cd24a1f3..90b26855 100644 --- a/pgdog/src/frontend/client/query_engine/context.rs +++ b/pgdog/src/frontend/client/query_engine/context.rs @@ -2,11 +2,12 @@ use crate::{ backend::pool::{connection::mirror::Mirror, stats::MemoryStats}, frontend::{ Client, ClientRequest, PreparedStatements, - client::{Sticky, TransactionType, timeouts::Timeouts}, + client::{ClientRequestSettings, Sticky, TransactionType}, router::parser::rewrite::statement::plan::RewriteResult, }, net::{FrontendPid, Parameters, Stream}, }; +use std::net::SocketAddr; /// Context passed to the query engine to execute a query. pub struct QueryEngineContext<'a> { @@ -24,8 +25,8 @@ pub struct QueryEngineContext<'a> { pub(super) stream: &'a mut Stream, /// Client in transaction? pub(super) transaction: Option, - /// Timeouts - pub(super) timeouts: Timeouts, + /// Per-request settings snapshot. + pub(super) request_settings: ClientRequestSettings, /// Cross shard queries are disabled. pub(super) cross_shard_disabled: Option, /// Client memory usage. @@ -38,10 +39,8 @@ pub struct QueryEngineContext<'a> { pub(super) sticky: Sticky, /// Rewrite result. pub(super) rewrite_result: Option, - /// Log queries to stdout. - pub(super) query_log_stdout: bool, - /// Maximum query message size before a warning is logged. - pub(super) query_size_limit: Option, + /// Client TCP address, used for `application_name_add_host`. + pub(super) client_addr: SocketAddr, } impl<'a> QueryEngineContext<'a> { @@ -55,7 +54,7 @@ impl<'a> QueryEngineContext<'a> { client_request: &mut client.client_request, stream: &mut client.stream, transaction: client.transaction, - timeouts: client.timeouts, + request_settings: client.request_settings, cross_shard_disabled: None, memory_stats, admin: client.admin, @@ -63,8 +62,7 @@ impl<'a> QueryEngineContext<'a> { rollback: false, sticky: client.sticky, rewrite_result: None, - query_log_stdout: client.query_log_stdout, - query_size_limit: client.query_size_limit, + client_addr: client.addr, } } @@ -83,7 +81,10 @@ impl<'a> QueryEngineContext<'a> { client_request: buffer, stream: &mut mirror.stream, transaction: mirror.transaction, - timeouts: mirror.timeouts, + request_settings: ClientRequestSettings { + timeouts: mirror.timeouts, + ..ClientRequestSettings::default() + }, cross_shard_disabled: None, memory_stats: MemoryStats::default(), admin: false, @@ -91,8 +92,7 @@ impl<'a> QueryEngineContext<'a> { rollback: false, sticky: Sticky::new(), rewrite_result: None, - query_log_stdout: false, - query_size_limit: None, + client_addr: SocketAddr::from(([0, 0, 0, 0], 0)), } } diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 484ae538..a65d92d4 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -1,6 +1,5 @@ use crate::{ backend::pool::{Connection, Request}, - config::config, frontend::{ BufferedQuery, Client, ClientComms, Command, Error, Router, RouterContext, Stats, client::query_engine::{hooks::QueryEngineHooks, route_query::ClusterCheck}, @@ -171,7 +170,7 @@ impl QueryEngine { .route // Admin commands don't have a route. .as_mut() .and_then(|route| route.take_explain()) - && config().config.general.expanded_explain + && context.request_settings.expanded_explain { self.pending_explain = Some(ExplainResponseState::new(trace)); } diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index e133706a..067b69ac 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -9,7 +9,6 @@ use crate::{ DataRow, FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, RowDescription, ToBytes, TransactionState, }, - state::State, util::safe_timeout, }; @@ -58,7 +57,10 @@ impl QueryEngine { } } - let query_timeout = context.timeouts.query_timeout(&State::Active); + let query_timeout = context + .request_settings + .timeouts + .query_timeout(&self.stats.state); let result = safe_timeout(query_timeout, self.client_server_exchange(context)).await; match result { @@ -320,9 +322,16 @@ impl QueryEngine { // Update client params with values // sent from the server using ParameterStatus(B) messages. if !changed_params.is_empty() { + let add_host = context.request_settings.application_name_add_host; + let host = context.client_addr.to_string(); for (name, value) in changed_params.iter() { + let value = if add_host && name.eq_ignore_ascii_case("application_name") { + value.clone().with_client_host(&host) + } else { + value.clone() + }; debug!("setting client's \"{}\" to {}", name, value); - context.params.insert(name.clone(), value.clone()); + context.params.insert(name.clone(), value); } self.comms.update_params(context.params); } diff --git a/pgdog/src/frontend/client/query_engine/query_log_stdout.rs b/pgdog/src/frontend/client/query_engine/query_log_stdout.rs index 5a943671..cb38eb07 100644 --- a/pgdog/src/frontend/client/query_engine/query_log_stdout.rs +++ b/pgdog/src/frontend/client/query_engine/query_log_stdout.rs @@ -1,12 +1,11 @@ use tracing::{info, warn}; use super::QueryEngineContext; -use crate::config::config; use crate::net::ProtocolMessage; use crate::util::{sanitize_log_sample, user_database_from_params}; pub(super) fn log_query_stdout(context: &QueryEngineContext<'_>) { - let size_limit = context.query_size_limit; + let size_limit = context.request_settings.query_size_limit; // Largest query message in the request, when it exceeds the limit. // The limit protects the query parser, so only messages carrying SQL @@ -22,7 +21,7 @@ pub(super) fn log_query_stdout(context: &QueryEngineContext<'_>) { .filter(|&size| size > size_limit) }); - if !context.query_log_stdout && oversize.is_none() { + if !context.request_settings.query_log_stdout && oversize.is_none() { return; } @@ -32,7 +31,7 @@ pub(super) fn log_query_stdout(context: &QueryEngineContext<'_>) { let one_line = sanitize_log_sample( query.query(), - config().config.general.log_query_sample_length, + context.request_settings.log_query_sample_length, ); let one_line = one_line.trim(); @@ -45,7 +44,7 @@ pub(super) fn log_query_stdout(context: &QueryEngineContext<'_>) { "[large_query] size={}B query_size_limit={}B '{}...' [database: {}, user: {}]", size, size_limit, one_line, database, user, ); - } else if context.query_log_stdout { + } else if context.request_settings.query_log_stdout { info!("{} [database: {}, user: {}]", one_line, database, user); } } diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 2e54a9a0..6ac5772f 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -55,7 +55,10 @@ impl QueryEngine { // Wait for boot-time maintenance before we throw traffic at the cluster. if let Ok(cluster) = self.backend.cluster() { safe_timeout( - context.timeouts.query_timeout(&State::Active), + context + .request_settings + .timeouts + .query_timeout(&State::Active), cluster.wait_ready(), ) .await diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index dff7cfde..bb327fef 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -1,5 +1,7 @@ +use crate::frontend::ClientRequest; use crate::frontend::SetParam; use crate::frontend::router::parameter_hints::{PGDOG_PIN, PGDOG_SHARD, PGDOG_SHARDING_KEY}; +use crate::net::ProtocolMessage; use crate::net::messages::ErrorResponse; use super::*; @@ -22,8 +24,23 @@ impl QueryEngine { return Ok(()); } + let mut params = params.to_vec(); + if context.request_settings.application_name_add_host { + let host = context.client_addr.to_string(); + for param in &mut params { + if param.name.eq_ignore_ascii_case("application_name") + && let Some(value) = param.value.take() + { + param.value = Some(value.with_client_host(&host)); + } + } + if !behave_like_select { + rewrite_application_name_query(context.client_request, ¶ms); + } + } + let mut fake_command = "SET"; - for param in params { + for param in ¶ms { let is_pin = param.name == PGDOG_PIN; if let Some(value) = param.value.clone() { @@ -108,3 +125,132 @@ impl QueryEngine { Ok(()) } } + +fn rewrite_application_name_query(request: &mut ClientRequest, params: &[SetParam]) { + let [param] = params else { + return; + }; + if !param.name.eq_ignore_ascii_case("application_name") { + return; + } + let Some(value) = ¶m.value else { + return; + }; + + let cmd = if param.local { "SET LOCAL" } else { "SET" }; + let sql = format!(r#"{cmd} "application_name" TO {value}"#); + for message in &mut request.messages { + if let ProtocolMessage::Query(query) = message { + query.set_query(&sql); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::frontend::SetParam; + use crate::net::Query; + use crate::net::parameter::ParameterValue; + + fn query_sql(request: &ClientRequest) -> &str { + match &request.messages[0] { + ProtocolMessage::Query(query) => query.query(), + other => panic!("expected Query, got {other:?}"), + } + } + + #[test] + fn rewrite_application_name_query_rewrites_set() { + let mut request = ClientRequest::default(); + request.push(ProtocolMessage::Query(Query::new( + "SET application_name TO 'client'", + ))); + let params = vec![SetParam { + name: "application_name".into(), + value: Some(ParameterValue::String("client - 127.0.0.1:1234".into())), + local: false, + }]; + + rewrite_application_name_query(&mut request, ¶ms); + + assert_eq!( + query_sql(&request), + r#"SET "application_name" TO "client - 127.0.0.1:1234""# + ); + } + + #[test] + fn rewrite_application_name_query_uses_set_local() { + let mut request = ClientRequest::default(); + request.push(ProtocolMessage::Query(Query::new( + "SET LOCAL application_name TO 'client'", + ))); + let params = vec![SetParam { + name: "application_name".into(), + value: Some(ParameterValue::String("client - 127.0.0.1:1234".into())), + local: true, + }]; + + rewrite_application_name_query(&mut request, ¶ms); + + assert_eq!( + query_sql(&request), + r#"SET LOCAL "application_name" TO "client - 127.0.0.1:1234""# + ); + } + + #[test] + fn rewrite_application_name_query_skips_non_application_name() { + let mut request = ClientRequest::default(); + request.push(ProtocolMessage::Query(Query::new("SET timezone TO 'UTC'"))); + let params = vec![SetParam { + name: "timezone".into(), + value: Some(ParameterValue::String("UTC".into())), + local: false, + }]; + + rewrite_application_name_query(&mut request, ¶ms); + + assert_eq!(query_sql(&request), "SET timezone TO 'UTC'"); + } + + #[test] + fn rewrite_application_name_query_skips_multi_param() { + let mut request = ClientRequest::default(); + request.push(ProtocolMessage::Query(Query::new( + "SET application_name TO 'client'", + ))); + let params = vec![ + SetParam { + name: "application_name".into(), + value: Some(ParameterValue::String("client".into())), + local: false, + }, + SetParam { + name: "timezone".into(), + value: Some(ParameterValue::String("UTC".into())), + local: false, + }, + ]; + + rewrite_application_name_query(&mut request, ¶ms); + + assert_eq!(query_sql(&request), "SET application_name TO 'client'"); + } + + #[test] + fn rewrite_application_name_query_skips_reset() { + let mut request = ClientRequest::default(); + request.push(ProtocolMessage::Query(Query::new("RESET application_name"))); + let params = vec![SetParam { + name: "application_name".into(), + value: None, + local: false, + }]; + + rewrite_application_name_query(&mut request, ¶ms); + + assert_eq!(query_sql(&request), "RESET application_name"); + } +} diff --git a/pgdog/src/frontend/client/query_engine/test/set.rs b/pgdog/src/frontend/client/query_engine/test/set.rs index 734ef4c4..305ffa75 100644 --- a/pgdog/src/frontend/client/query_engine/test/set.rs +++ b/pgdog/src/frontend/client/query_engine/test/set.rs @@ -1,11 +1,14 @@ use crate::{ backend::databases::reload_from_existing, - config::{config, load_test_sharded, set}, + config::{config, load_test, load_test_sharded, set}, expect_message, - net::{CommandComplete, ErrorResponse, ReadyForQuery, parameter::ParameterValue}, + net::{ + CommandComplete, DataRow, ErrorResponse, Format, ReadyForQuery, RowDescription, + parameter::ParameterValue, + }, }; -use super::prelude::*; +use super::{change_config, prelude::*}; /// Number of shards the client is currently connected to. fn connected_servers(client: &mut TestClient) -> usize { @@ -731,3 +734,112 @@ async fn test_lock_timeout() { "lock_timeout should be cleared after RESET" ); } + +#[tokio::test] +async fn test_set_application_name_add_host() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + change_config(|g| g.application_name_add_host = true); + + test_client + .send_simple(Query::new("SET application_name TO 'toto'")) + .await; + + assert_eq!( + expect_message!(test_client.read().await, CommandComplete).command(), + "SET" + ); + assert_eq!( + expect_message!(test_client.read().await, ReadyForQuery).status, + 'I' + ); + + assert_eq!( + test_client.client().params.get("application_name").unwrap(), + &ParameterValue::String("toto - 127.0.0.1:1234".into()), + "SET application_name should re-append the client host" + ); + + change_config(|g| g.application_name_add_host = false); +} + +#[tokio::test] +async fn test_set_application_name_add_host_with_backend() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + change_config(|g| g.application_name_add_host = true); + + test_client.send_simple(Query::new("BEGIN")).await; + test_client.read_until('Z').await.unwrap(); + + test_client.send_simple(Query::new("SELECT 1")).await; + test_client.read_until('Z').await.unwrap(); + assert!(test_client.backend_connected()); + + test_client + .send_simple(Query::new("SET application_name TO 'server'")) + .await; + test_client.read_until('Z').await.unwrap(); + + assert_eq!( + test_client.client().params.get("application_name").unwrap(), + &ParameterValue::String("server - 127.0.0.1:1234".into()), + "backend SET should sync suffixed application_name from ParameterStatus" + ); + + test_client.send_simple(Query::new("ROLLBACK")).await; + test_client.read_until('Z').await.unwrap(); + + change_config(|g| g.application_name_add_host = false); +} + +#[tokio::test] +async fn test_set_config_application_name_add_host() { + let mut test_client = TestClient::new_sharded(Parameters::default()).await; + change_config(|g| g.application_name_add_host = true); + + test_client + .send_simple(Query::new( + "SELECT set_config('application_name', 'cfg', false)", + )) + .await; + + expect_message!(test_client.read().await, RowDescription); + let row = expect_message!(test_client.read().await, DataRow); + assert_eq!( + row.get::(0, Format::Text).unwrap(), + "cfg - 127.0.0.1:1234" + ); + expect_message!(test_client.read().await, CommandComplete); + expect_message!(test_client.read().await, ReadyForQuery); + + assert_eq!( + test_client.client().params.get("application_name").unwrap(), + &ParameterValue::String("cfg - 127.0.0.1:1234".into()), + ); + + change_config(|g| g.application_name_add_host = false); +} + +#[tokio::test] +async fn test_startup_application_name_add_host() { + load_test(); + change_config(|g| g.application_name_add_host = true); + + let mut params = Parameters::default(); + params.insert("application_name", "startup"); + let client = Client::new_test(Stream::dev_null(), params); + + assert_eq!( + client.params.get("application_name").unwrap(), + &ParameterValue::String("startup - 127.0.0.1:1234".into()), + "startup application_name should include client host" + ); + + let client = Client::new_test(Stream::dev_null(), Parameters::default()); + assert_eq!( + client.params.get("application_name").unwrap(), + &ParameterValue::String(" - 127.0.0.1:1234".into()), + "missing application_name should keep an empty prefix" + ); + + change_config(|g| g.application_name_add_host = false); +} diff --git a/pgdog/src/frontend/client/request_settings.rs b/pgdog/src/frontend/client/request_settings.rs new file mode 100644 index 00000000..12bcdf06 --- /dev/null +++ b/pgdog/src/frontend/client/request_settings.rs @@ -0,0 +1,143 @@ +use pgdog_config::General; + +use super::timeouts::Timeouts; + +/// Per-request snapshot of client settings read from config. +/// +/// Filled once in [`super::Client::buffer`] and passed through the query engine +/// so mid-request config reloads don't change behavior. +#[derive(Debug, Clone, Copy)] +pub(crate) struct ClientRequestSettings { + pub(crate) timeouts: Timeouts, + pub(crate) query_log_stdout: bool, + pub(crate) query_size_limit: Option, + pub(crate) application_name_add_host: bool, + pub(crate) expanded_explain: bool, + pub(crate) log_query_sample_length: usize, + pub(crate) frontend_query_size_limit_block: Option, +} + +impl Default for ClientRequestSettings { + fn default() -> Self { + Self { + timeouts: Timeouts::default(), + query_log_stdout: false, + query_size_limit: None, + application_name_add_host: false, + expanded_explain: false, + log_query_sample_length: General::log_query_sample_length(), + frontend_query_size_limit_block: None, + } + } +} + +impl ClientRequestSettings { + pub(crate) fn from_general(general: &General) -> Self { + Self { + timeouts: Timeouts::from_config(general), + query_log_stdout: general.query_log_stdout, + query_size_limit: general.query_size_limit, + application_name_add_host: general.application_name_add_host, + expanded_explain: general.expanded_explain, + log_query_sample_length: general.log_query_sample_length, + frontend_query_size_limit_block: general.frontend_query_size_limit_block(), + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pgdog_config::QuerySizeLimitAction; + + use crate::{config::General, frontend::ClientRequest, state::State}; + + use super::*; + + #[test] + fn default_settings_match_safe_defaults() { + let settings = ClientRequestSettings::default(); + + assert!(!settings.query_log_stdout); + assert_eq!(settings.query_size_limit, None); + assert!(!settings.application_name_add_host); + assert!(!settings.expanded_explain); + assert_eq!( + settings.log_query_sample_length, + General::log_query_sample_length() + ); + assert_eq!(settings.frontend_query_size_limit_block, None); + assert_eq!( + settings.timeouts.query_timeout(&State::Active), + Duration::MAX + ); + assert_eq!( + settings + .timeouts + .client_idle_timeout(&State::Idle, &ClientRequest::default()), + Duration::MAX + ); + assert_eq!( + settings + .timeouts + .client_idle_timeout(&State::IdleInTransaction, &ClientRequest::default()), + Duration::MAX + ); + } + + #[test] + fn from_general_copies_all_snapshotted_fields() { + let general = General { + query_timeout: 1_000, + client_idle_timeout: 2_000, + client_idle_in_transaction_timeout: 3_000, + query_log_stdout: true, + query_size_limit: Some(4096), + query_size_limit_action: QuerySizeLimitAction::Block, + application_name_add_host: true, + expanded_explain: true, + log_query_sample_length: 42, + ..Default::default() + }; + + let settings = ClientRequestSettings::from_general(&general); + + assert!(settings.query_log_stdout); + assert_eq!(settings.query_size_limit, Some(4096)); + assert!(settings.application_name_add_host); + assert!(settings.expanded_explain); + assert_eq!(settings.log_query_sample_length, 42); + assert_eq!(settings.frontend_query_size_limit_block, Some(4096)); + assert_eq!( + settings.timeouts.query_timeout(&State::Active), + Duration::from_millis(1_000) + ); + assert_eq!( + settings + .timeouts + .client_idle_timeout(&State::Idle, &ClientRequest::default()), + Duration::from_millis(2_000) + ); + assert_eq!( + settings + .timeouts + .client_idle_timeout(&State::IdleInTransaction, &ClientRequest::default()), + Duration::from_millis(3_000) + ); + } + + #[test] + fn from_general_omits_block_limit_when_action_is_warn() { + let general = General { + query_size_limit: Some(1024), + query_size_limit_action: QuerySizeLimitAction::Warn, + ..Default::default() + }; + + let settings = ClientRequestSettings::from_general(&general); + + assert_eq!(settings.query_size_limit, Some(1024)); + assert_eq!(settings.frontend_query_size_limit_block, None); + } +} diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 7d290976..3a6294db 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -173,6 +173,26 @@ impl ParameterValue { _ => None, } } + + /// Append ` - {host}` for `application_name_add_host`. Idempotent if the + /// suffix is already present. + pub fn with_client_host(self, host: &str) -> Self { + match self { + Self::String(name) => Self::String(application_name_with_host(&name, host)), + other => other, + } + } +} + +/// Format `application_name` with the client host, matching +/// `{name} - {host}` (empty `{name}` is allowed). +pub fn application_name_with_host(name: &str, host: &str) -> String { + let suffix = format!(" - {host}"); + if name.ends_with(&suffix) { + name.to_string() + } else { + format!("{name}{suffix}") + } } /// List of parameters. @@ -548,7 +568,33 @@ mod test { use crate::net::ToBytes; use crate::net::parameter::ParameterValue; - use super::Parameters; + use super::{Parameters, application_name_with_host}; + + #[test] + fn test_application_name_with_host() { + assert_eq!( + application_name_with_host("myapp", "10.0.0.5:1234"), + "myapp - 10.0.0.5:1234" + ); + assert_eq!( + application_name_with_host("", "10.0.0.5:1234"), + " - 10.0.0.5:1234" + ); + assert_eq!( + application_name_with_host("myapp - 10.0.0.5:1234", "10.0.0.5:1234"), + "myapp - 10.0.0.5:1234" + ); + assert_eq!( + application_name_with_host("psql", "[::1]:54321"), + "psql - [::1]:54321" + ); + } + + #[test] + fn test_with_client_host_non_string_is_unchanged() { + let value = ParameterValue::Integer(1).with_client_host("127.0.0.1:1234"); + assert_eq!(value, ParameterValue::Integer(1)); + } #[test] fn test_identical() {