Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Escape PostgreSQL options #3800

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
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
22 changes: 20 additions & 2 deletions sqlx-postgres/src/options/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::borrow::Cow;
use std::env::var;
use std::fmt::{Display, Write};
use std::fmt::Display;
use std::path::{Path, PathBuf};

pub use ssl_mode::PgSslMode;
Expand Down Expand Up @@ -495,6 +495,9 @@ impl PgConnectOptions {

/// Set additional startup options for the connection as a list of key-value pairs.
///
/// Escapes the options’ backslash and space characters as per
/// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
///
/// # Example
///
/// ```rust
Expand All @@ -515,7 +518,16 @@ impl PgConnectOptions {
options_str.push(' ');
}

write!(options_str, "-c {k}={v}").expect("failed to write an option to the string");
// Escape options per
// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
options_str.push_str("-c ");
for c in format!("{k}={v}").chars() {
match c {
'\\' => options_str.push_str("\\\\"),
' ' => options_str.push_str("\\ "),
c => options_str.push(c),
};
}
Comment on lines +524 to +530
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This bothers me for a few reasons:

  • Formats the string and escapes it afterward
  • Allocates a whole string just to throw it away
  • Pushes to options_str character-by-character
    • This especially pessimizes the case where no escaping is needed.

The latter two both serve to create unnecessary allocation pressure. While this isn't a super hot routine, a lot of little allocations tend to increase heap fragmentation, which is compounded by the number of options being passed to a connection, the lengths of the values, and the number of connections that are opened over the lifetime of the application.

With just an extra ten lines of code, it's possible to avoid the intermediate allocation and write the whole string at once if there are no characters to escape: https://play.rust-lang.org/?version=stable&mode=debug&edition=2024&gist=0b0bb2909484f9578c2e4c02ed1f41b5

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You would use it like this:

use std::fmt::Write;

// Writing to a string is infallible.
write!(options_str, "-c {}={}", PgEscapeOption(k), PgEscapeOption(v)).ok();

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the quick and thorough feedback! I like your suggestion and tried to apply it – but:

The key and value parameters of the options function are of the Display trait. To avoid heap allocations your solution would have to be adopted to struct PgEscapeOption<T: Display>(T);. I don’t think that is possible.

One can sidestep the problem by changing the parameters of the options function to take a string instead, so their backing buffer can be accessed. However, that would slightly restrict the user in what they can pass in.

pub fn options<'a, I>(mut self, options: I) -> Self
    where
        I: IntoIterator<Item = (&'a str, &'a str)>,
    {

Formats the string and escapes it afterward

I don’t think that does any harm on its own. I agree with the other points.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The key and value parameters of the options function are of the Display trait. To avoid heap allocations your solution would have to be adopted to struct PgEscapeOption<T: Display>(T);. I don’t think that is possible.

I realized that only late last night. It is possible, you just need to drop down a layer: https://play.rust-lang.org/?version=stable&mode=debug&edition=2024&gist=a8a59849ca9c7d1f3a76d4bfe1ab05f8

Not quite as elegant but still avoids the intermediate allocations.

I don’t think that does any harm on its own. I agree with the other points.

No, it's just a little weird.

}
self
}
Expand Down Expand Up @@ -683,6 +695,12 @@ fn test_options_formatting() {
options.options,
Some("-c geqo=off -c statement_timeout=5min".to_string())
);
// https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
let options = PgConnectOptions::new().options([("passfile", r"/back\slash/ and\ spaces")]);
assert_eq!(
options.options,
Some(r"-c passfile=/back\\slash/\ and\\\ spaces".to_string())
);
let options = PgConnectOptions::new();
assert_eq!(options.options, None);
}