add mysql cancel token

Author: IsaacHorvath

src/lib.rs

@@ -128,6 +128,9 @@ mod transaction_manager;
 #[cfg(feature = "mysql")]
 #[doc(inline)]
 pub use self::mysql::AsyncMysqlConnection;
+#[cfg(feature = "mysql")]
+#[doc(inline)]
+pub use self::mysql::CancelToken;
 #[cfg(feature = "postgres")]
 #[doc(inline)]
 pub use self::pg::AsyncPgConnection;

src/mysql/cancel_token.rs

@@ -0,0 +1,37 @@
+use mysql_async::prelude::Query;
+use mysql_async::{Opts, OptsBuilder};
+
+use crate::mysql::error_helper::ErrorHelper;
+
+/// The capability to request cancellation of in-progress queries on a
+/// connection.
+#[derive(Clone)]
+pub struct CancelToken {
+    pub(crate) opts: Opts,
+    pub(crate) kill_id: u32,
+}
+
+impl CancelToken {
+    /// Attempts to cancel the in-progress query on the connection associated
+    /// with this `CancelToken`.
+    ///
+    /// The server provides no information about whether a cancellation attempt was successful or not. An error will
+    /// only be returned if the client was unable to connect to the database.
+    ///
+    /// Cancellation is inherently racy. There is no guarantee that the
+    /// cancellation request will reach the server before the query terminates
+    /// normally, or that the connection associated with this token is still
+    /// active.
+    pub async fn cancel_query(&self) -> Result<(), diesel::result::ConnectionError> {
+        let builder = OptsBuilder::from_opts(self.opts.clone());
+
+        let conn = mysql_async::Conn::new(builder).await.map_err(ErrorHelper)?;
+
+        format!("KILL QUERY {};", self.kill_id)
+            .ignore(conn)
+            .await
+            .map_err(ErrorHelper)?;
+
+        Ok(())
+    }
+}

src/mysql/mod.rs

@@ -19,10 +19,12 @@ use mysql_async::prelude::Queryable;
 use mysql_async::{Opts, OptsBuilder, Statement};
 use std::future::Future;
 
+mod cancel_token;
 mod error_helper;
 mod row;
 mod serialize;
 
+pub use self::cancel_token::CancelToken;
 use self::error_helper::ErrorHelper;
 use self::row::MysqlRow;
 use self::serialize::ToSqlHelper;
@@ -254,6 +256,14 @@ impl AsyncMysqlConnection {
         Ok(conn)
     }
 
+    /// Constructs a cancellation token that can later be used to request cancellation of a query running on the connection associated with this client.
+    pub fn cancel_token(&self) -> CancelToken {
+        let kill_id = self.conn.id();
+        let opts = self.conn.opts().clone();
+
+        CancelToken { kill_id, opts }
+    }
+
     fn with_prepared_statement<'conn, T, F, R>(
         &'conn mut self,
         query: T,

tests/lib.rs

@@ -167,6 +167,51 @@ async fn postgres_cancel_token() {
     }
 }
 
+#[cfg(feature = "mysql")]
+#[tokio::test]
+async fn mysql_cancel_token() {
+    use std::time::Duration;
+
+    use diesel::result::{DatabaseErrorKind, Error};
+
+    let (sender, receiver) = tokio::sync::oneshot::channel();
+
+    // execute a long-running query on a separate thread
+    let task = tokio::spawn(async move {
+        let conn = &mut connection().await;
+        let token = conn.cancel_token();
+
+        // send the token back to the main thread via a oneshot channel
+        sender
+            .send(token)
+            .unwrap_or_else(|_| panic!("couldn't send token"));
+
+        diesel::dsl::sql::<diesel::sql_types::Integer>("SELECT SLEEP(5)")
+            .load::<i32>(conn)
+            .await
+    });
+
+    // wait for the cancellation token to be sent
+    if let Ok(token) = receiver.await {
+        // give the query time to start before invoking the token
+        tokio::time::sleep(Duration::from_millis(500)).await;
+        token.cancel_query().await.unwrap();
+    }
+
+    // make sure the query task resulted in a cancellation error or a return value of 1:
+    match task.await.unwrap() {
+        Err(e) => match e {
+            Error::DatabaseError(DatabaseErrorKind::Unknown, v)
+                if v.message() == "Query execution was interrupted" => {}
+            _ => panic!("unexpected error: {:?}", e),
+        },
+        Ok(r) => match r[0] {
+            1 => {}
+            _ => panic!(""),
+        },
+    }
+}
+
 #[cfg(feature = "postgres")]
 async fn setup(connection: &mut TestConnection) {
     diesel::sql_query(