fix(postgres): syntax error in EXPLAIN query

Author: abonander

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 {

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);
     }
 }

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);
@@ -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)
     }
 }
 
@@ -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)
     }
 }
 
@@ -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);
@@ -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);
+}