From 37f53cc7e92421fe5277399ec6219b898047a556 Mon Sep 17 00:00:00 2001
From: Austin Bonander <austin.bonander@gmail.com>
Date: Wed, 21 Aug 2024 14:49:21 -0700
Subject: [PATCH] fix(postgres): syntax error in EXPLAIN query

---
 sqlx-postgres/src/connection/describe.rs |  6 ++-
 sqlx-postgres/src/io/buf_mut.rs          | 10 +---
 sqlx-postgres/src/io/mod.rs              | 64 +++++++++++++++++-------
 3 files changed, 52 insertions(+), 28 deletions(-)

diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs
index 145616fec1..9a46a202d5 100644
--- a/sqlx-postgres/src/connection/describe.rs
+++ b/sqlx-postgres/src/connection/describe.rs
@@ -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 {
diff --git a/sqlx-postgres/src/io/buf_mut.rs b/sqlx-postgres/src/io/buf_mut.rs
index eea9d34acd..0fe3809b57 100644
--- a/sqlx-postgres/src/io/buf_mut.rs
+++ b/sqlx-postgres/src/io/buf_mut.rs
@@ -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);
     }
 }
diff --git a/sqlx-postgres/src/io/mod.rs b/sqlx-postgres/src/io/mod.rs
index df064b1e2c..72f2a978c8 100644
--- a/sqlx-postgres/src/io/mod.rs
+++ b/sqlx-postgres/src/io/mod.rs
@@ -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);
@@ -35,16 +40,22 @@ impl StatementId {
         self.0.name_len(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)
+    /// 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)
+    }
+
+    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)
     }
 }
 
@@ -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)
     }
 }
 
@@ -93,6 +104,11 @@ impl IdInner {
         )
     }
 
+    #[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);
@@ -113,18 +129,28 @@ impl IdInner {
     }
 
     #[inline(always)]
-    fn write_name<E>(
-        &self,
-        name_prefix: &str,
-        mut write: impl FnMut(&str) -> Result<(), E>,
-    ) -> Result<(), E> {
+    fn put_name_with_nul(&self, name_prefix: &str, 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);
+}