Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .schema/pgdog.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
"description": "General settings are relevant to the operations of the pooler itself, or apply to all database pools.\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/>",
"$ref": "#/$defs/General",
"default": {
"application_name_add_host": false,
"auth_type": "scram",
"ban_replica_lag": 9223372036854775807,
"ban_replica_lag_bytes": 9223372036854775807,
Expand Down Expand Up @@ -653,6 +654,11 @@
"description": "General settings are relevant to the operations of the pooler itself, or apply to all database pools.\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/>",
"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<https://docs.pgdog.dev/configuration/pgdog.toml/general/#application_name_add_host>",
"type": "boolean",
"default": false
},
"auth_type": {
"description": "What kind of authentication mechanism to use for client connections.\n\n_Default:_ `scram`\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/#auth_type>",
"$ref": "#/$defs/AuthType",
Expand Down
22 changes: 22 additions & 0 deletions pgdog-config/src/general.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Comment thread
jkaczman marked this conversation as resolved.
///
/// 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`
///
/// <https://docs.pgdog.dev/configuration/pgdog.toml/general/#application_name_add_host>
#[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.
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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]
Expand Down
63 changes: 40 additions & 23 deletions pgdog/src/frontend/client/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand All @@ -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;

Expand Down Expand Up @@ -73,10 +77,8 @@ pub struct Client {
prepared_statements: PreparedStatements,
// Client transaction state.
transaction: Option<TransactionType>,
// 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.
Expand All @@ -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<usize>,
}

/// Inputs to the per-user client certificate check.
Expand Down Expand Up @@ -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<ConfigAndUsers>,
protocol_version: ProtocolVersion,
Expand All @@ -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(&params);
let admin = database == config.config.admin.name && config.config.admin.user == user;
let admin_password = &config.config.admin.password;
Expand Down Expand Up @@ -420,19 +423,29 @@ 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,
config.config.general.frontend_query_size_limit_block(),
),
sticky: Sticky::from_params(&params),
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;
Expand All @@ -442,21 +455,28 @@ 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();
prepared_statements.level = config().config.general.prepared_statements;

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,
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -736,7 +753,7 @@ impl MemoryUsage for Client {
+ std::mem::size_of::<ClientComms>()
+ std::mem::size_of::<bool>() * 5
+ self.prepared_statements.memory_used()
+ std::mem::size_of::<Timeouts>()
+ std::mem::size_of::<ClientRequestSettings>()
+ self.stream_buffer.capacity()
+ self.client_request.memory_usage()
}
Expand Down
5 changes: 4 additions & 1 deletion pgdog/src/frontend/client/query_engine/connect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
26 changes: 13 additions & 13 deletions pgdog/src/frontend/client/query_engine/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand All @@ -24,8 +25,8 @@ pub struct QueryEngineContext<'a> {
pub(super) stream: &'a mut Stream,
/// Client in transaction?
pub(super) transaction: Option<TransactionType>,
/// 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<bool>,
/// Client memory usage.
Expand All @@ -38,10 +39,8 @@ pub struct QueryEngineContext<'a> {
pub(super) sticky: Sticky,
/// Rewrite result.
pub(super) rewrite_result: Option<RewriteResult>,
/// 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<usize>,
/// Client TCP address, used for `application_name_add_host`.
pub(super) client_addr: SocketAddr,
}

impl<'a> QueryEngineContext<'a> {
Expand All @@ -55,16 +54,15 @@ 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,
requests_left: 0,
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,
}
}

Expand All @@ -83,16 +81,18 @@ 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,
requests_left: 0,
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)),
}
}

Expand Down
3 changes: 1 addition & 2 deletions pgdog/src/frontend/client/query_engine/mod.rs
Original file line number Diff line number Diff line change
@@ -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},
Expand Down Expand Up @@ -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));
}
Expand Down
Loading
Loading