diff --git a/libs/pq_proto/src/lib.rs b/libs/pq_proto/src/lib.rs index c5e4dbd1f0..b7995c840c 100644 --- a/libs/pq_proto/src/lib.rs +++ b/libs/pq_proto/src/lib.rs @@ -75,27 +75,36 @@ impl StartupMessageParams { /// taking into account all escape sequences but leaving them as-is. /// [`None`] means that there's no `options` in [`Self`]. pub fn options_raw(&self) -> Option> { - // See `postgres: pg_split_opts`. - let mut last_was_escape = false; - let iter = self - .get("options")? - .split(move |c: char| { - // We split by non-escaped whitespace symbols. - let should_split = c.is_ascii_whitespace() && !last_was_escape; - last_was_escape = c == '\\' && !last_was_escape; - should_split - }) - .filter(|s| !s.is_empty()); - - Some(iter) + self.get("options").map(Self::parse_options_raw) } /// Split command-line options according to PostgreSQL's logic, /// applying all escape sequences (using owned strings as needed). /// [`None`] means that there's no `options` in [`Self`]. pub fn options_escaped(&self) -> Option>> { + self.get("options").map(Self::parse_options_escaped) + } + + /// Split command-line options according to PostgreSQL's logic, + /// taking into account all escape sequences but leaving them as-is. + pub fn parse_options_raw(input: &str) -> impl Iterator { // See `postgres: pg_split_opts`. - let iter = self.options_raw()?.map(|s| { + let mut last_was_escape = false; + input + .split(move |c: char| { + // We split by non-escaped whitespace symbols. + let should_split = c.is_ascii_whitespace() && !last_was_escape; + last_was_escape = c == '\\' && !last_was_escape; + should_split + }) + .filter(|s| !s.is_empty()) + } + + /// Split command-line options according to PostgreSQL's logic, + /// applying all escape sequences (using owned strings as needed). + pub fn parse_options_escaped(input: &str) -> impl Iterator> { + // See `postgres: pg_split_opts`. + Self::parse_options_raw(input).map(|s| { let mut preserve_next_escape = false; let escape = |c| { // We should remove '\\' unless it's preceded by '\\'. @@ -108,9 +117,12 @@ impl StartupMessageParams { true => Cow::Owned(s.replace(escape, "")), false => Cow::Borrowed(s), } - }); + }) + } - Some(iter) + /// Iterate through key-value pairs in an arbitrary order. + pub fn iter(&self) -> impl Iterator { + self.params.iter().map(|(k, v)| (k.as_str(), v.as_str())) } // This function is mostly useful in tests. diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index 3f5eb3caff..2e12d9ee26 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -77,14 +77,9 @@ impl ConnCfg { self.dbname(dbname); } - if let Some(options) = params.options_raw() { - // We must drop all proxy-specific parameters. - #[allow(unstable_name_collisions)] - let options: String = options - .filter(|opt| !opt.starts_with("project=")) - .intersperse(" ") // TODO: use impl from std once it's stabilized - .collect(); - + // 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); } @@ -225,3 +220,46 @@ impl ConnCfg { Ok(connection) } } + +/// Retrieve `options` from a startup message, dropping all proxy-secific flags. +fn filtered_options(params: &StartupMessageParams) -> Option { + #[allow(unstable_name_collisions)] + let options: String = params + .options_raw()? + .filter(|opt| !opt.starts_with("project=")) + .intersperse(" ") // TODO: use impl from std once it's stabilized + .collect(); + + // Don't even bother with empty options. + if options.is_empty() { + return None; + } + + Some(options) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_filtered_options() { + // Empty options is unlikely to be useful anyway. + let params = StartupMessageParams::new([("options", "")]); + assert_eq!(filtered_options(¶ms), 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); + + // Same, because unescaped whitespaces are no-op. + let params = StartupMessageParams::new([("options", " project=foo ")]); + assert_eq!(filtered_options(¶ms).as_deref(), None); + + let params = StartupMessageParams::new([("options", r"\ project=foo \ ")]); + assert_eq!(filtered_options(¶ms).as_deref(), Some(r"\ \ ")); + + let params = StartupMessageParams::new([("options", "project = foo")]); + assert_eq!(filtered_options(¶ms).as_deref(), Some("project = foo")); + } +}