diff --git a/Cargo.lock b/Cargo.lock index 77bf01240273..70c837c14645 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4005,7 +4005,7 @@ dependencies = [ [[package]] name = "postgres" version = "0.19.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#20031d7a9ee1addeae6e0968e3899ae6bf01cee2" +source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#cff6927e4f58b1af6ecc2ee7279df1f2ff537295" dependencies = [ "bytes", "fallible-iterator", @@ -4018,7 +4018,7 @@ dependencies = [ [[package]] name = "postgres-protocol" version = "0.6.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#20031d7a9ee1addeae6e0968e3899ae6bf01cee2" +source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#cff6927e4f58b1af6ecc2ee7279df1f2ff537295" dependencies = [ "base64 0.20.0", "byteorder", @@ -4037,7 +4037,7 @@ dependencies = [ [[package]] name = "postgres-types" version = "0.2.4" -source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#20031d7a9ee1addeae6e0968e3899ae6bf01cee2" +source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#cff6927e4f58b1af6ecc2ee7279df1f2ff537295" dependencies = [ "bytes", "fallible-iterator", @@ -6210,7 +6210,7 @@ dependencies = [ [[package]] name = "tokio-postgres" version = "0.7.7" -source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#20031d7a9ee1addeae6e0968e3899ae6bf01cee2" +source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#cff6927e4f58b1af6ecc2ee7279df1f2ff537295" dependencies = [ "async-trait", "byteorder", diff --git a/libs/postgres_connection/src/lib.rs b/libs/postgres_connection/src/lib.rs index 9f57f3d50750..fdabcbacb245 100644 --- a/libs/postgres_connection/src/lib.rs +++ b/libs/postgres_connection/src/lib.rs @@ -144,20 +144,7 @@ impl PgConnectionConfig { // implement and this function is hardly a bottleneck. The function is only called around // establishing a new connection. #[allow(unstable_name_collisions)] - config.options( - &self - .options - .iter() - .map(|s| { - if s.contains(['\\', ' ']) { - Cow::Owned(s.replace('\\', "\\\\").replace(' ', "\\ ")) - } else { - Cow::Borrowed(s.as_str()) - } - }) - .intersperse(Cow::Borrowed(" ")) // TODO: use impl from std once it's stabilized - .collect::(), - ); + config.options(&encode_options(&self.options)); } config } @@ -178,6 +165,21 @@ impl PgConnectionConfig { } } +#[allow(unstable_name_collisions)] +fn encode_options(options: &[String]) -> String { + options + .iter() + .map(|s| { + if s.contains(['\\', ' ']) { + Cow::Owned(s.replace('\\', "\\\\").replace(' ', "\\ ")) + } else { + Cow::Borrowed(s.as_str()) + } + }) + .intersperse(Cow::Borrowed(" ")) // TODO: use impl from std once it's stabilized + .collect::() +} + impl fmt::Display for PgConnectionConfig { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // The password is intentionally hidden and not part of this display string. @@ -206,7 +208,7 @@ impl fmt::Debug for PgConnectionConfig { #[cfg(test)] mod tests_pg_connection_config { - use crate::PgConnectionConfig; + use crate::{encode_options, PgConnectionConfig}; use once_cell::sync::Lazy; use url::Host; @@ -255,18 +257,12 @@ mod tests_pg_connection_config { #[test] fn test_with_options() { - let cfg = PgConnectionConfig::new_host_port(STUB_HOST.clone(), 123).extend_options([ - "hello", - "world", - "with space", - "and \\ backslashes", + let options = encode_options(&[ + "hello".to_owned(), + "world".to_owned(), + "with space".to_owned(), + "and \\ backslashes".to_owned(), ]); - assert_eq!(cfg.host(), &*STUB_HOST); - assert_eq!(cfg.port(), 123); - assert_eq!(cfg.raw_address(), "stub.host.example:123"); - assert_eq!( - cfg.to_tokio_postgres_config().get_options(), - Some("hello world with\\ space and\\ \\\\\\ backslashes") - ); + assert_eq!(options, "hello world with\\ space and\\ \\\\\\ backslashes"); } } diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index feb09d563896..a50a96e5e844 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -103,12 +103,8 @@ impl ConnCfg { /// Reuse password or auth keys from the other config. pub fn reuse_password(&mut self, other: Self) { - if let Some(password) = other.get_password() { - self.password(password); - } - - if let Some(keys) = other.get_auth_keys() { - self.auth_keys(keys); + if let Some(password) = other.get_auth() { + self.auth(password); } } @@ -124,48 +120,64 @@ impl ConnCfg { /// Apply startup message params to the connection config. pub fn set_startup_params(&mut self, params: &StartupMessageParams) { - // Only set `user` if it's not present in the config. - // Link auth flow takes username from the console's response. - if let (None, Some(user)) = (self.get_user(), params.get("user")) { - self.user(user); - } - - // Only set `dbname` if it's not present in the config. - // Link auth flow takes dbname from the console's response. - if let (None, Some(dbname)) = (self.get_dbname(), params.get("database")) { - self.dbname(dbname); - } - - // Don't add `options` if they were only used for specifying a project. - // Connection pools don't support `options`, because they affect backend startup. - if let Some(options) = filtered_options(params) { - self.options(&options); - } - - if let Some(app_name) = params.get("application_name") { - self.application_name(app_name); - } - - // TODO: This is especially ugly... - if let Some(replication) = params.get("replication") { - use tokio_postgres::config::ReplicationMode; - match replication { - "true" | "on" | "yes" | "1" => { - self.replication_mode(ReplicationMode::Physical); + let mut client_encoding = false; + for (k, v) in params.iter() { + match k { + "user" => { + // Only set `user` if it's not present in the config. + // Link auth flow takes username from the console's response. + if self.get_user().is_none() { + self.user(v); + } } "database" => { - self.replication_mode(ReplicationMode::Logical); + // Only set `dbname` if it's not present in the config. + // Link auth flow takes dbname from the console's response. + if self.get_dbname().is_none() { + self.dbname(v); + } + } + "options" => { + // Don't add `options` if they were only used for specifying a project. + // Connection pools don't support `options`, because they affect backend startup. + if let Some(options) = filtered_options(v) { + self.options(&options); + } + } + + // the special ones in tokio-postgres that we don't want being set by the user + "dbname" => {} + "password" => {} + "sslmode" => {} + "host" => {} + "port" => {} + "connect_timeout" => {} + "keepalives" => {} + "keepalives_idle" => {} + "keepalives_interval" => {} + "keepalives_retries" => {} + "target_session_attrs" => {} + "channel_binding" => {} + "max_backend_message_size" => {} + + "client_encoding" => { + client_encoding = true; + // only error should be from bad null bytes, + // but we've already checked for those. + _ = self.param("client_encoding", v); + } + + _ => { + // only error should be from bad null bytes, + // but we've already checked for those. + _ = self.param(k, v); } - _other => {} } } - - // TODO: extend the list of the forwarded startup parameters. - // Currently, tokio-postgres doesn't allow us to pass - // arbitrary parameters, but the ones above are a good start. - // - // This and the reverse params problem can be better addressed - // in a bespoke connection machinery (a new library for that sake). + if !client_encoding { + // for compatibility since we removed it from tokio-postgres + self.param("client_encoding", "UTF8").unwrap(); + } } } @@ -338,10 +350,9 @@ impl ConnCfg { } /// Retrieve `options` from a startup message, dropping all proxy-secific flags. -fn filtered_options(params: &StartupMessageParams) -> Option { +fn filtered_options(options: &str) -> Option { #[allow(unstable_name_collisions)] - let options: String = params - .options_raw()? + let options: String = StartupMessageParams::parse_options_raw(options) .filter(|opt| parse_endpoint_param(opt).is_none() && neon_option(opt).is_none()) .intersperse(" ") // TODO: use impl from std once it's stabilized .collect(); @@ -413,27 +424,23 @@ mod tests { #[test] fn test_filtered_options() { // Empty options is unlikely to be useful anyway. - let params = StartupMessageParams::new([("options", "")]); - assert_eq!(filtered_options(¶ms), None); + assert_eq!(filtered_options(""), None); // It's likely that clients will only use options to specify endpoint/project. - let params = StartupMessageParams::new([("options", "project=foo")]); - assert_eq!(filtered_options(¶ms), None); + let params = "project=foo"; + assert_eq!(filtered_options(params), None); // Same, because unescaped whitespaces are no-op. - let params = StartupMessageParams::new([("options", " project=foo ")]); - assert_eq!(filtered_options(¶ms).as_deref(), None); + let params = " project=foo "; + assert_eq!(filtered_options(params), None); - let params = StartupMessageParams::new([("options", r"\ project=foo \ ")]); - assert_eq!(filtered_options(¶ms).as_deref(), Some(r"\ \ ")); + let params = r"\ project=foo \ "; + assert_eq!(filtered_options(params).as_deref(), Some(r"\ \ ")); - let params = StartupMessageParams::new([("options", "project = foo")]); - assert_eq!(filtered_options(¶ms).as_deref(), Some("project = foo")); + let params = "project = foo"; + assert_eq!(filtered_options(params).as_deref(), Some("project = foo")); - let params = StartupMessageParams::new([( - "options", - "project = foo neon_endpoint_type:read_write neon_lsn:0/2", - )]); - assert_eq!(filtered_options(¶ms).as_deref(), Some("project = foo")); + let params = "project = foo neon_endpoint_type:read_write neon_lsn:0/2"; + assert_eq!(filtered_options(params).as_deref(), Some("project = foo")); } } diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 86e64c0a386d..05d60612385c 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -231,6 +231,10 @@ impl ConnectMechanism for TokioMechanism { .dbname(&self.conn_info.dbname) .connect_timeout(timeout); + config + .param("client_encoding", "UTF8") + .expect("client encoding UTF8 is always valid"); + let pause = ctx.latency_timer.pause(crate::metrics::Waiting::Compute); let res = config.connect(tokio_postgres::NoTls).await; drop(pause); diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 7a99aeb75938..583ff75f7ca7 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -202,6 +202,7 @@ fn get_conn_info( options = Some(NeonOptions::parse_options_raw(&value)); } } + ctx.set_db_options(params.freeze()); let user_info = ComputeUserInfo { endpoint, diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index f446f4f200d7..8ed44b109442 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -53,6 +53,25 @@ def test_proxy_select_1(static_proxy: NeonProxy): assert out[0][0] == 42 +def test_proxy_server_params(static_proxy: NeonProxy): + """ + Test that server params are passing through to postgres + """ + + out = static_proxy.safe_psql( + "select to_json('0 seconds'::interval)", options="-c intervalstyle=iso_8601" + ) + assert out[0][0] == "PT0S" + out = static_proxy.safe_psql( + "select to_json('0 seconds'::interval)", options="-c intervalstyle=sql_standard" + ) + assert out[0][0] == "0" + out = static_proxy.safe_psql( + "select to_json('0 seconds'::interval)", options="-c intervalstyle=postgres" + ) + assert out[0][0] == "00:00:00" + + def test_password_hack(static_proxy: NeonProxy): """ Check the PasswordHack auth flow: an alternative to SCRAM auth for