Merge pull request #204 from blackbeam/query-byte-slice

Accept either bytes or strings for queryes (updated)

Cargo.toml

@@ -22,7 +22,7 @@ futures-sink = "0.3"
 lazy_static = "1"
 lru = "0.7.0"
 mio = { version = "0.8.0", features = ["os-poll", "net"] }
-mysql_common = { version = "0.28.0", default-features = false }
+mysql_common = { version = "0.29.0", default-features = false }
 native-tls = "0.2"
 once_cell = "1.7.2"
 pem = "1.0.1"

src/conn/mod.rs

@@ -1422,8 +1422,8 @@ mod test {
             .stmt_cache_ref()
             .iter()
             .map(|item| item.1.query.0.as_ref())
-            .collect::<Vec<&str>>();
-        assert_eq!(order, &["DO 6", "DO 5", "DO 3"]);
+            .collect::<Vec<&[u8]>>();
+        assert_eq!(order, &[b"DO 6", b"DO 5", b"DO 3"]);
         conn.disconnect().await?;
         Ok(())
     }

src/conn/pool/mod.rs

@@ -388,7 +388,7 @@ mod test {
 
             // now we'll kill connections..
             for id in ids {
-                master.query_drop(&format!("KILL {}", id)).await?;
+                master.query_drop(format!("KILL {}", id)).await?;
             }
 
             // now check, that they're still in the pool..

src/conn/routines/prepare.rs

@@ -11,21 +11,21 @@ use super::Routine;
 /// A routine that performs `COM_STMT_PREPARE`.
 #[derive(Debug, Clone)]
 pub struct PrepareRoutine {
-    query: Arc<str>,
+    query: Arc<[u8]>,
 }
 
 impl PrepareRoutine {
-    pub fn new(raw_query: Cow<'_, str>) -> Self {
+    pub fn new(raw_query: Cow<'_, [u8]>) -> Self {
         Self {
-            query: raw_query.into_owned().into_boxed_str().into(),
+            query: raw_query.into_owned().into_boxed_slice().into(),
         }
     }
 }
 
 impl Routine<Arc<StmtInner>> for PrepareRoutine {
     fn call<'a>(&'a mut self, conn: &'a mut Conn) -> BoxFuture<'a, crate::Result<Arc<StmtInner>>> {
         async move {
-            conn.write_command_data(Command::COM_STMT_PREPARE, self.query.as_bytes())
+            conn.write_command_data(Command::COM_STMT_PREPARE, &self.query)
                 .await?;
 
             let packet = conn.read_packet().await?;

src/conn/stmt_cache.rs

@@ -19,16 +19,16 @@ use std::{
 use crate::queryable::stmt::StmtInner;
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct QueryString(pub Arc<str>);
+pub struct QueryString(pub Arc<[u8]>);
 
-impl Borrow<str> for QueryString {
-    fn borrow(&self) -> &str {
+impl Borrow<[u8]> for QueryString {
+    fn borrow(&self) -> &[u8] {
         &*self.0.as_ref()
     }
 }
 
-impl PartialEq<str> for QueryString {
-    fn eq(&self, other: &str) -> bool {
+impl PartialEq<[u8]> for QueryString {
+    fn eq(&self, other: &[u8]) -> bool {
         &*self.0.as_ref() == other
     }
 }
@@ -68,7 +68,7 @@ impl StmtCache {
         }
     }
 
-    pub fn put(&mut self, query: Arc<str>, stmt: Arc<StmtInner>) -> Option<Arc<StmtInner>> {
+    pub fn put(&mut self, query: Arc<[u8]>, stmt: Arc<StmtInner>) -> Option<Arc<StmtInner>> {
         if self.cap == 0 {
             return None;
         }
@@ -95,7 +95,7 @@ impl StmtCache {
 
     pub fn remove(&mut self, id: u32) {
         if let Some(entry) = self.cache.pop(&id) {
-            self.query_map.remove::<str>(entry.query.borrow());
+            self.query_map.remove::<[u8]>(entry.query.borrow());
         }
     }
 
@@ -135,7 +135,7 @@ impl super::Conn {
     /// Returns statement, if cached.
     ///
     /// `raw_query` is the query with `?` placeholders (not with `:<name>` placeholders).
-    pub(crate) fn get_cached_stmt(&mut self, raw_query: &str) -> Option<Arc<StmtInner>> {
+    pub(crate) fn get_cached_stmt(&mut self, raw_query: &[u8]) -> Option<Arc<StmtInner>> {
         self.stmt_cache_mut()
             .by_query(raw_query)
             .map(|entry| entry.stmt.clone())

src/error.rs

@@ -107,8 +107,8 @@ pub enum DriverError {
     #[error("Error converting from mysql row.")]
     FromRow { row: Row },
 
-    #[error("Missing named parameter `{}'.", name)]
-    MissingNamedParam { name: String },
+    #[error("Missing named parameter `{}'.", String::from_utf8_lossy(&name))]
+    MissingNamedParam { name: Vec<u8> },
 
     #[error("Named and positional parameters mixed in one statement.")]
     MixedParams,

src/lib.rs

@@ -404,6 +404,8 @@ pub mod prelude {
     #[doc(inline)]
     pub use crate::local_infile_handler::GlobalHandler;
     #[doc(inline)]
+    pub use crate::query::AsQuery;
+    #[doc(inline)]
     pub use crate::query::{BatchQuery, Query, WithParams};
     #[doc(inline)]
     pub use crate::queryable::Queryable;

src/query.rs

@@ -6,6 +6,8 @@
 // option. All files in the project carrying such notice may not be copied,
 // modified, or distributed except according to those terms.
 
+use std::borrow::Cow;
+
 use futures_util::FutureExt;
 
 use crate::{
@@ -15,6 +17,54 @@ use crate::{
     BinaryProtocol, BoxFuture, Params, QueryResult, ResultSetStream, TextProtocol,
 };
 
+/// Types that can be treated as a MySQL query.
+///
+/// This trait is implemented by all "string-ish" standard library types, like `String`, `&str`,
+/// `Cow<str>`, but also all types that can be treated as a slice of bytes (such as `Vec<u8>` and
+/// `&[u8]`), since MySQL does not require queries to be valid UTF-8.
+pub trait AsQuery: Send + Sync {
+    fn as_query(&self) -> Cow<'_, [u8]>;
+}
+
+impl AsQuery for &'_ [u8] {
+    fn as_query(&self) -> Cow<'_, [u8]> {
+        Cow::Borrowed(self)
+    }
+}
+
+macro_rules! impl_as_query_as_ref {
+    ($type: ty) => {
+        impl AsQuery for $type {
+            fn as_query(&self) -> Cow<'_, [u8]> {
+                Cow::Borrowed(self.as_ref())
+            }
+        }
+    };
+}
+
+impl_as_query_as_ref!(Vec<u8>);
+impl_as_query_as_ref!(&Vec<u8>);
+impl_as_query_as_ref!(Box<[u8]>);
+impl_as_query_as_ref!(Cow<'_, [u8]>);
+impl_as_query_as_ref!(std::sync::Arc<[u8]>);
+
+macro_rules! impl_as_query_as_bytes {
+    ($type: ty) => {
+        impl AsQuery for $type {
+            fn as_query(&self) -> Cow<'_, [u8]> {
+                Cow::Borrowed(self.as_bytes())
+            }
+        }
+    };
+}
+
+impl_as_query_as_bytes!(String);
+impl_as_query_as_bytes!(&String);
+impl_as_query_as_bytes!(&str);
+impl_as_query_as_bytes!(Box<str>);
+impl_as_query_as_bytes!(Cow<'_, str>);
+impl_as_query_as_bytes!(std::sync::Arc<str>);
+
 /// MySql text query.
 ///
 /// This trait covers the set of `query*` methods on the `Queryable` trait.
@@ -157,7 +207,7 @@ pub trait Query: Send + Sized {
     }
 }
 
-impl<Q: AsRef<str> + Send + Sync> Query for Q {
+impl<Q: AsQuery> Query for Q {
     type Protocol = TextProtocol;
 
     fn run<'a, 't: 'a, C>(self, conn: C) -> BoxFuture<'a, QueryResult<'a, 't, TextProtocol>>

src/queryable/mod.rs

@@ -28,6 +28,7 @@ use crate::{
     consts::CapabilityFlags,
     error::*,
     prelude::{FromRow, StatementLike},
+    query::AsQuery,
     queryable::query_result::ResultSetMeta,
     BoxFuture, Column, Conn, Params, ResultSetStream, Row,
 };
@@ -102,9 +103,9 @@ impl Conn {
     /// Low level function that performs a text query.
     pub(crate) async fn raw_query<'a, Q>(&'a mut self, query: Q) -> Result<()>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
     {
-        self.routine(QueryRoutine::new(query.as_ref().as_bytes()))
+        self.routine(QueryRoutine::new(query.as_query().as_ref()))
             .await
     }
 }
@@ -122,7 +123,7 @@ pub trait Queryable: Send {
         query: Q,
     ) -> BoxFuture<'a, QueryResult<'a, 'static, TextProtocol>>
     where
-        Q: AsRef<str> + Send + Sync + 'a;
+        Q: AsQuery + 'a;
 
     /// Prepares the given statement.
     ///
@@ -165,7 +166,7 @@ pub trait Queryable: Send {
     /// to make this conversion infallible.
     fn query<'a, T, Q>(&'a mut self, query: Q) -> BoxFuture<'a, Vec<T>>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
         T: FromRow + Send + 'static,
     {
         async move { self.query_iter(query).await?.collect_and_drop::<T>().await }.boxed()
@@ -180,7 +181,7 @@ pub trait Queryable: Send {
     /// to make this conversion infallible.
     fn query_first<'a, T, Q>(&'a mut self, query: Q) -> BoxFuture<'a, Option<T>>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
         T: FromRow + Send + 'static,
     {
         async move {
@@ -205,7 +206,7 @@ pub trait Queryable: Send {
     /// to make this conversion infallible.
     fn query_map<'a, T, F, Q, U>(&'a mut self, query: Q, mut f: F) -> BoxFuture<'a, Vec<U>>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
         T: FromRow + Send + 'static,
         F: FnMut(T) -> U + Send + 'a,
         U: Send,
@@ -229,7 +230,7 @@ pub trait Queryable: Send {
     /// to make this conversion infallible.
     fn query_fold<'a, T, F, Q, U>(&'a mut self, query: Q, init: U, mut f: F) -> BoxFuture<'a, U>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
         T: FromRow + Send + 'static,
         F: FnMut(U, T) -> U + Send + 'a,
         U: Send + 'a,
@@ -246,7 +247,7 @@ pub trait Queryable: Send {
     /// Performs the given query and drops the query result.
     fn query_drop<'a, Q>(&'a mut self, query: Q) -> BoxFuture<'a, ()>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
     {
         async move { self.query_iter(query).await?.drop_result().await }.boxed()
     }
@@ -397,7 +398,7 @@ pub trait Queryable: Send {
     ) -> BoxFuture<'a, ResultSetStream<'a, 'a, 'static, T, TextProtocol>>
     where
         T: Unpin + FromRow + Send + 'static,
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
     {
         async move {
             self.query_iter(query)
@@ -451,10 +452,10 @@ impl Queryable for Conn {
         query: Q,
     ) -> BoxFuture<'a, QueryResult<'a, 'static, TextProtocol>>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
     {
         async move {
-            self.routine(QueryRoutine::new(query.as_ref().as_bytes()))
+            self.routine(QueryRoutine::new(query.as_query().as_ref()))
                 .await?;
             Ok(QueryResult::new(self))
         }
@@ -525,7 +526,7 @@ impl Queryable for Transaction<'_> {
         query: Q,
     ) -> BoxFuture<'a, QueryResult<'a, 'static, TextProtocol>>
     where
-        Q: AsRef<str> + Send + Sync + 'a,
+        Q: AsQuery + 'a,
     {
         self.0.query_iter(query)
     }

src/queryable/stmt.rs

@@ -22,6 +22,8 @@ use crate::{
     Column, Params,
 };
 
+use super::AsQuery;
+
 /// Result of a `StatementLike::to_statement` call.
 pub enum ToStatementResult<'a> {
     /// Statement is immediately available.
@@ -37,12 +39,13 @@ pub trait StatementLike: Send + Sync {
         Self: 'a;
 }
 
-fn to_statement_move<'a, T: AsRef<str> + Send + Sync + 'a>(
+fn to_statement_move<'a, T: AsQuery + 'a>(
     stmt: T,
     conn: &'a mut crate::Conn,
 ) -> ToStatementResult<'a> {
     let fut = async move {
-        let (named_params, raw_query) = parse_named_params(stmt.as_ref())?;
+        let query = stmt.as_query();
+        let (named_params, raw_query) = parse_named_params(query.as_ref())?;
         let inner_stmt = match conn.get_cached_stmt(&*raw_query) {
             Some(inner_stmt) => inner_stmt,
             None => conn.prepare_statement(raw_query).await?,
@@ -119,7 +122,7 @@ impl<T: StatementLike + Clone> StatementLike for &'_ T {
 /// Statement data.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct StmtInner {
-    pub(crate) raw_query: Arc<str>,
+    pub(crate) raw_query: Arc<[u8]>,
     columns: Option<Box<[Column]>>,
     params: Option<Box<[Column]>>,
     stmt_packet: StmtPacket,
@@ -130,7 +133,7 @@ impl StmtInner {
     pub(crate) fn from_payload(
         pld: &[u8],
         connection_id: u32,
-        raw_query: Arc<str>,
+        raw_query: Arc<[u8]>,
     ) -> std::io::Result<Self> {
         let stmt_packet = ParseBuf(pld).parse(())?;
 
@@ -192,11 +195,11 @@ impl StmtInner {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Statement {
     pub(crate) inner: Arc<StmtInner>,
-    pub(crate) named_params: Option<Vec<String>>,
+    pub(crate) named_params: Option<Vec<Vec<u8>>>,
 }
 
 impl Statement {
-    pub(crate) fn new(inner: Arc<StmtInner>, named_params: Option<Vec<String>>) -> Self {
+    pub(crate) fn new(inner: Arc<StmtInner>, named_params: Option<Vec<Vec<u8>>>) -> Self {
         Self {
             inner,
             named_params,
@@ -275,7 +278,7 @@ impl crate::Conn {
     /// Low-level helper, that prepares the given statement.
     ///
     /// `raw_query` is a query with `?` placeholders (if any).
-    async fn prepare_statement(&mut self, raw_query: Cow<'_, str>) -> Result<Arc<StmtInner>> {
+    async fn prepare_statement(&mut self, raw_query: Cow<'_, [u8]>) -> Result<Arc<StmtInner>> {
         let inner_stmt = self.routine(PrepareRoutine::new(raw_query)).await?;
 
         if let Some(old_stmt) = self.cache_stmt(&inner_stmt) {