Skip to content

Commit

Permalink
fix(postgres): syntax error in EXPLAIN query
Browse files Browse the repository at this point in the history
  • Loading branch information
abonander committed Aug 21, 2024
1 parent e731cfd commit 1efef02
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 27 deletions.
6 changes: 5 additions & 1 deletion sqlx-postgres/src/connection/describe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -500,7 +500,11 @@ WHERE rngtypid = $1
stmt_id: StatementId,
params_len: usize,
) -> Result<Vec<Option<bool>>, Error> {
let mut explain = format!("EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE {stmt_id}");
let stmt_id_display = stmt_id
.display()
.ok_or_else(|| err_protocol!("cannot EXPLAIN unnamed statement: {stmt_id:?}"))?;

let mut explain = format!("EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE {stmt_id_display}");
let mut comma = false;

if params_len > 0 {
Expand Down
10 changes: 2 additions & 8 deletions sqlx-postgres/src/io/buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,18 +47,12 @@ impl PgBufMutExt for Vec<u8> {
// writes a statement name by ID
#[inline]
fn put_statement_name(&mut self, id: StatementId) {
let _: Result<(), ()> = id.write_name(|s| {
self.extend_from_slice(s.as_bytes());
Ok(())
});
id.put_name_with_nul(self);
}

// writes a portal name by ID
#[inline]
fn put_portal_name(&mut self, id: PortalId) {
let _: Result<(), ()> = id.write_name(|s| {
self.extend_from_slice(s.as_bytes());
Ok(())
});
id.put_name_with_nul(self);
}
}
69 changes: 51 additions & 18 deletions sqlx-postgres/src/io/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ pub(crate) struct PortalId(IdInner);
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct IdInner(Option<NonZeroU32>);

pub(crate) struct DisplayId {
prefix: &'static str,
id: NonZeroU32,
}

impl StatementId {
#[allow(dead_code)]
pub const UNNAMED: Self = Self(IdInner::UNNAMED);
Expand All @@ -34,17 +39,23 @@ impl StatementId {
pub fn name_len(&self) -> Saturating<usize> {
self.0.name_len(Self::NAME_PREFIX)
}

/// Get a type to format this statement ID with [`Display`].
///
/// Returns `None` if this is the unnamed statement.
#[inline(always)]
pub fn display(&self) -> Option<DisplayId> {
self.0.display(Self::NAME_PREFIX)
}

// There's no common trait implemented by `Formatter` and `Vec<u8>` for this purpose;
// we're deliberately avoiding the formatting machinery because it's known to be slow.
pub fn write_name<E>(&self, write: impl FnMut(&str) -> Result<(), E>) -> Result<(), E> {
self.0.write_name(Self::NAME_PREFIX, write)
pub fn put_name_with_nul(&self, buf: &mut Vec<u8>) {
self.0.put_name_with_nul(Self::NAME_PREFIX, buf)
}
}

impl Display for StatementId {
impl Display for DisplayId {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
self.write_name(|s| f.write_str(s))
write!(f, "{}{}", self.prefix, self.id)
}
}

Expand All @@ -67,13 +78,13 @@ impl PortalId {
Self(self.0.next())
}

/// Calculate the number of bytes that will be written by [`Self::write_name()`].
/// Calculate the number of bytes that will be written by [`Self::put_name_with_nul()`].
pub fn name_len(&self) -> Saturating<usize> {
self.0.name_len(Self::NAME_PREFIX)
}

pub fn write_name<E>(&self, write: impl FnMut(&str) -> Result<(), E>) -> Result<(), E> {
self.0.write_name(Self::NAME_PREFIX, write)
pub fn put_name_with_nul(&self, buf: &mut Vec<u8>) {
self.0.put_name_with_nul(Self::NAME_PREFIX, buf)
}
}

Expand All @@ -92,7 +103,15 @@ impl IdInner {
.map(|id| id.checked_add(1).unwrap_or(NonZeroU32::MIN)),
)
}


#[inline(always)]
fn display(&self, prefix: &'static str) -> Option<DisplayId> {
self.0.map(|id| DisplayId {
prefix,
id
})
}

#[inline(always)]
fn name_len(&self, name_prefix: &str) -> Saturating<usize> {
let mut len = Saturating(0);
Expand All @@ -113,18 +132,32 @@ impl IdInner {
}

#[inline(always)]
fn write_name<E>(
fn put_name_with_nul(
&self,
name_prefix: &str,
mut write: impl FnMut(&str) -> Result<(), E>,
) -> Result<(), E> {
buf: &mut Vec<u8>,
) {
if let Some(id) = self.0 {
write(name_prefix)?;
write(itoa::Buffer::new().format(id.get()))?;
buf.extend_from_slice(name_prefix.as_bytes());
buf.extend_from_slice(itoa::Buffer::new().format(id.get()).as_bytes());
}

write("\0")?;

Ok(())
buf.push(0);
}
}

#[test]
fn statement_id_display_matches_encoding() {
const EXPECTED_STR: &str = "sqlx_s_1234567890";
const EXPECTED_BYTES: &[u8] = b"sqlx_s_1234567890\0";

let mut bytes = Vec::new();

StatementId::TEST_VAL.put_name_with_nul(&mut bytes);

assert_eq!(bytes, EXPECTED_BYTES);

let str = StatementId::TEST_VAL.display().unwrap().to_string();

assert_eq!(str, EXPECTED_STR);
}

0 comments on commit 1efef02

Please sign in to comment.