fix(cubesql): Query cancellation for simple query protocol (#5987)

Author: ovr

rust/cubesql/Cargo.lock

@@ -60,7 +60,7 @@ pretty_assertions = "1.0.0"
 insta = "1.12"
 mysql_async = "0.29"
 portpicker = "0.1.1"
-tokio-postgres = { version = "0.7.6", features = ["with-chrono-0_4", "runtime"] }
+tokio-postgres = { version = "0.7.7", features = ["with-chrono-0_4", "runtime"] }
 rust_decimal = { version = "1.23", features = ["db-tokio-postgres"] }
 pg_interval = "0.4.1"
 criterion = { version = "0.4.0", features = ["html_reports"] }

rust/cubesql/cubesql/Cargo.toml

@@ -264,22 +264,59 @@ impl PostgresIntegrationTestSuite {
         }
     }
 
-    async fn test_cancel(&self) -> RunResult<()> {
-        let cancel_token = self.client.cancel_token();
+    async fn test_cancel_execute_prepared(&self) -> RunResult<()> {
+        let client = PostgresIntegrationTestSuite::create_client(
+            format!("host=127.0.0.1 port={} user=test password=test", self.port)
+                .parse()
+                .unwrap(),
+        )
+        .await;
+
+        let cancel_token = client.cancel_token();
         let cancel = async move {
-            tokio::time::sleep(Duration::from_millis(1000)).await;
+            sleep(Duration::from_millis(1000)).await;
 
             cancel_token.cancel_query(NoTls).await
         };
 
         // testing_blocking tables will neven finish. It's a special testing table
-        let sleep = self
-            .client
-            .batch_execute("SELECT * FROM information_schema.testing_blocking");
+        let sleep = client.batch_execute("SELECT * FROM information_schema.testing_blocking");
+
+        match join!(sleep, cancel) {
+            (Err(ref e), Ok(())) if e.code() == Some(&SqlState::QUERY_CANCELED) => {}
+            res => panic!(
+                "unexpected return, prepared must be cancelled, actual: {:?}",
+                res
+            ),
+        };
+
+        Ok(())
+    }
+
+    async fn test_cancel_simple_query(&self) -> RunResult<()> {
+        let client = PostgresIntegrationTestSuite::create_client(
+            format!("host=127.0.0.1 port={} user=test password=test", self.port)
+                .parse()
+                .unwrap(),
+        )
+        .await;
+
+        let cancel_token = client.cancel_token();
+        let cancel = async move {
+            sleep(Duration::from_millis(1000)).await;
+
+            cancel_token.cancel_query(NoTls).await
+        };
+
+        // testing_blocking tables will neven finish. It's a special testing table
+        let sleep = client.simple_query("SELECT * FROM information_schema.testing_blocking");
 
         match join!(sleep, cancel) {
             (Err(ref e), Ok(())) if e.code() == Some(&SqlState::QUERY_CANCELED) => {}
-            t => panic!("unexpected return {:?}", t),
+            (_, err) => panic!(
+                "unexpected return, simple query must be cancelled, actual: {:?}",
+                err
+            ),
         };
 
         Ok(())
@@ -830,7 +867,8 @@ impl AsyncTestSuite for PostgresIntegrationTestSuite {
     }
 
     async fn run(&mut self) -> RunResult<()> {
-        self.test_cancel().await?;
+        self.test_cancel_simple_query().await?;
+        self.test_cancel_execute_prepared().await?;
         self.test_prepare().await?;
         self.test_extended_error().await?;
         self.test_prepare_empty_query().await?;

rust/cubesql/cubesql/e2e/tests/postgres.rs

@@ -931,20 +931,27 @@ impl AsyncPostgresShim {
 
         tokio::select! {
             _ = cancel.cancelled() => {
-                if let Some(qtrace) = qtrace {
-                    qtrace.set_statement_error_message("Execution cancelled by user");
-                }
                 self.session.state.end_query();
 
                 // We don't return error, because query can contains multiple statements
                 // then cancel request will cancel only one query
                 self.write(protocol::ErrorResponse::query_canceled()).await?;
+                if let Some(qtrace) = qtrace {
+                    qtrace.set_statement_error_message("Execution cancelled by user");
+                }
 
                 Ok(())
             },
             res = self.process_simple_query(stmt, meta, cancel.clone(), qtrace) => {
                 self.session.state.end_query();
 
+                if cancel.is_cancelled() {
+                    self.write(protocol::ErrorResponse::query_canceled()).await?;
+                    if let Some(qtrace) = qtrace {
+                        qtrace.set_statement_error_message("Execution cancelled by user");
+                    }
+                }
+
                 res
             },
         }
@@ -1362,25 +1369,30 @@ impl AsyncPostgresShim {
         cancel: CancellationToken,
     ) -> Result<(), ConnectionError> {
         let mut writer = BatchWriter::new(portal.get_format());
-        let completion = portal.execute(&mut writer, max_rows).await?;
 
-        if cancel.is_cancelled() {
-            return Ok(());
-        }
-
-        // Special handling for special queries, such as DISCARD ALL.
-        if let Some(description) = portal.get_description()? {
-            match description.len() {
-                0 => self.write(protocol::NoData::new()).await?,
-                _ => self.write(description).await?,
-            };
-        }
+        tokio::select! {
+            _ = cancel.cancelled() => {
+                // TODO: Cancellation handling via errors?
+                return Ok(());
+            },
+            res = portal.execute(&mut writer, max_rows) => {
+                let completion = res?;
+
+                // Special handling for special queries, such as DISCARD ALL.
+                if let Some(description) = portal.get_description()? {
+                    match description.len() {
+                        0 => self.write(protocol::NoData::new()).await?,
+                        _ => self.write(description).await?,
+                    };
+                }
 
-        if writer.has_data() {
-            buffer::write_direct(&mut self.socket, writer).await?;
-        };
+                if writer.has_data() {
+                    buffer::write_direct(&mut self.socket, writer).await?;
+                };
 
-        self.write_completion(completion).await
+                self.write_completion(completion).await
+            }
+        }
     }
 
     /// Pipeline of Execution