test: prove mysql driver is cancellation safe

Cargo.lock

@@ -1123,6 +1123,7 @@ dependencies = [
  "sha-1",
  "sha2",
  "sqlx-core",
+ "sqlx-test",
  "tokio",
  "url",
 ]
@@ -1152,6 +1153,15 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "sqlx-test"
+version = "0.0.0"
+dependencies = [
+ "anyhow",
+ "futures-util",
+ "tokio",
+]
+
 [[package]]
 name = "subtle"
 version = "2.4.0"

Cargo.toml

@@ -1,9 +1,4 @@
 [workspace]
 default-members = ["sqlx"]
 
-members = [
-    "sqlx-core",
-    "sqlx-mysql",
-    "sqlx-postgres",
-    "sqlx",
-]
+members = ["sqlx-core", "sqlx-mysql", "sqlx-postgres", "sqlx-test", "sqlx"]

sqlx-mysql/Cargo.toml

@@ -43,6 +43,7 @@ rand = "0.7"
 
 [dev-dependencies]
 sqlx-core = { version = "0.6.0-pre", path = "../sqlx-core", features = ["_mock"] }
+sqlx-test = { path = "../sqlx-test" }
 futures-executor = "0.3.8"
 anyhow = "1.0.37"
 tokio = { version = "1.0", features = ["full"] }

sqlx-mysql/src/connection/flush.rs

@@ -44,7 +44,7 @@ macro_rules! impl_flush {
 
                 Command::Simple => {
                     // simple commands where we expect an OK or ERR
-                    // ex. COM_PING, COM_QUERY, COM_STMT_RESET, COM_SET_OPTION
+                    // ex. COM_PING, COM_STMT_RESET, COM_SET_OPTION
                     maybe_end_with(commands, read_packet!($(@$blocking)? stream).deserialize_with(capabilities)?);
                 }
 

sqlx-mysql/tests/connection.rs

@@ -1,9 +1,11 @@
-use sqlx_core::{Connect, Connection, Tokio};
+use sqlx_core::{Connect, Connection, Executor, Tokio};
+use sqlx_mysql::MySqlArguments;
 use sqlx_mysql::MySqlConnection;
+use sqlx_test::assert_cancellation_safe;
 use std::env;
 
 #[tokio::test]
-async fn it_connects() -> anyhow::Result<()> {
+async fn test_connect() -> anyhow::Result<()> {
     let url = env::var("DATABASE_URL")?;
     let mut conn = MySqlConnection::<Tokio>::connect(&url).await?;
 
@@ -11,3 +13,46 @@ async fn it_connects() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_select_1() -> anyhow::Result<()> {
+    let url = env::var("DATABASE_URL")?;
+    let mut conn = MySqlConnection::<Tokio>::connect(&url).await?;
+
+    let row = conn.fetch_one("SELECT 1").await?;
+    let col0: i32 = row.try_get(0)?;
+
+    assert_eq!(col0, 1);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_ping_cancel() -> anyhow::Result<()> {
+    let url = env::var("DATABASE_URL")?;
+    let mut conn = MySqlConnection::<Tokio>::connect(&url).await?;
+
+    assert_cancellation_safe(&mut conn, |conn| conn.ping(), |conn| conn.ping()).await?;
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_select_cancel() -> anyhow::Result<()> {
+    let url = env::var("DATABASE_URL")?;
+    let mut conn = MySqlConnection::<Tokio>::connect(&url).await?;
+
+    assert_cancellation_safe(
+        &mut conn,
+        |conn| {
+            let mut args = MySqlArguments::new();
+            args.add_unchecked(&1_i32);
+
+            conn.fetch_one(("SELECT ?", args))
+        },
+        |conn| conn.ping(),
+    )
+    .await?;
+
+    Ok(())
+}

sqlx-test/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "sqlx-test"
+version = "0.0.0"
+edition = "2018"
+publish = false
+description = "Internal crate providing test utilities for testing SQLx database drivers."
+
+[dependencies]
+tokio = { version = "1.0", features = ["full"] }
+futures-util = "0.3.8"
+anyhow = "1.0.37"

sqlx-test/src/cancellation.rs

@@ -0,0 +1,51 @@
+use futures_util::future::{poll_fn, BoxFuture};
+use futures_util::pin_mut;
+use std::error::Error as StdError;
+use std::future::Future;
+use std::task::Poll;
+
+pub async fn assert_cancellation_safe<C, F1, F1R, F1E, F2, F2R, F2E>(
+    mut context: C,
+    task: F1,
+    checkpoint: F2,
+) -> anyhow::Result<()>
+where
+    F1: Fn(&mut C) -> BoxFuture<Result<F1R, F1E>>,
+    F1E: 'static + StdError + Send + Sync,
+    F2: Fn(&mut C) -> BoxFuture<Result<F2R, F2E>>,
+    F2E: 'static + StdError + Send + Sync,
+{
+    for _ in 0..100 {
+        for max_polls in 0.. {
+            let mut num_polls = 0;
+
+            {
+                let fut = (task)(&mut context);
+                pin_mut!(fut);
+
+                let res = poll_fn(|ctx| {
+                    let poll = match fut.as_mut().poll(ctx) {
+                        Poll::Ready(it) => Poll::Ready(Some(it)),
+                        Poll::Pending if num_polls == max_polls => Poll::Ready(None),
+                        Poll::Pending => Poll::Pending,
+                    };
+
+                    num_polls += 1;
+
+                    poll
+                })
+                .await;
+
+                match res {
+                    Some(Ok(_)) => break,
+                    Some(Err(error)) => return Err(error.into()),
+                    None => {}
+                }
+            }
+
+            (checkpoint)(&mut context).await?;
+        }
+    }
+
+    Ok(())
+}

sqlx-test/src/lib.rs

@@ -0,0 +1,3 @@
+mod cancellation;
+
+pub use cancellation::assert_cancellation_safe;