Merge pull request #206 from estie-inc/feat/query-cancel

feat: add explicit query cancellation API

Author: yuly3

Cargo.toml

@@ -169,3 +169,9 @@ name = "statement_bind"
 required-features = [
     "external-browser-sso",
 ]
+
+[[example]]
+name = "query_cancel"
+required-features = [
+    "external-browser-sso",
+]

examples/query_cancel.rs

@@ -0,0 +1,70 @@
+use std::{env, time::Duration};
+
+use snowflake_connector_rs::{
+    AuthConfig, Client, ClientConfig, ExternalBrowserConfig, QueryCancelStatus, QueryOptions,
+    Session, SessionConfig,
+};
+
+#[tokio::main]
+async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
+    let session = build_session().await?;
+
+    let query = session.query_handle_with_options(
+        "CALL SYSTEM$WAIT(120)",
+        QueryOptions::new().with_query_cancel_request_timeout(Duration::from_secs(10)),
+    )?;
+
+    let canceller = query.canceller();
+    let execution = tokio::spawn(query.execute());
+
+    tokio::time::sleep(Duration::from_secs(1)).await;
+
+    let status = canceller.cancel().await?;
+    if status == QueryCancelStatus::NotSubmitted {
+        println!("cancelled before the query was submitted");
+    } else {
+        println!("cancel status: {status:?}");
+    }
+
+    let execution_result = execution.await?;
+    match execution_result {
+        Ok(_cursor) => println!("query completed before the cancel took effect"),
+        Err(error) if error.is_cancelled() => println!("query was cancelled: {error}"),
+        Err(error) => return Err(error.into()),
+    }
+
+    Ok(())
+}
+
+async fn build_session() -> std::result::Result<Session, Box<dyn std::error::Error>> {
+    let username = env::var("SNOWFLAKE_USERNAME")?;
+    let account = env::var("SNOWFLAKE_ACCOUNT")?;
+    let role = env::var("SNOWFLAKE_ROLE").ok();
+    let warehouse = env::var("SNOWFLAKE_WAREHOUSE").ok();
+    let database = env::var("SNOWFLAKE_DATABASE").ok();
+    let schema = env::var("SNOWFLAKE_SCHEMA").ok();
+
+    let mut session_config = SessionConfig::new();
+    if let Some(value) = warehouse {
+        session_config = session_config.with_warehouse(value);
+    }
+    if let Some(value) = database {
+        session_config = session_config.with_database(value);
+    }
+    if let Some(value) = schema {
+        session_config = session_config.with_schema(value);
+    }
+    if let Some(value) = role {
+        session_config = session_config.with_role(value);
+    }
+
+    let client = Client::new(
+        ClientConfig::new(
+            &username,
+            &account,
+            AuthConfig::external_browser(ExternalBrowserConfig::default()),
+        )
+        .with_session(session_config),
+    )?;
+    Ok(client.create_session().await?)
+}

src/config.rs

@@ -31,6 +31,7 @@ pub struct SessionConfig {
 }
 
 pub(crate) const DEFAULT_QUERY_RESPONSE_TIMEOUT: Duration = Duration::from_secs(300);
+pub(crate) const DEFAULT_QUERY_CANCEL_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
 const DEFAULT_COLLECT_PREFETCH_CONCURRENCY: NonZeroUsize =
     NonZeroUsize::new(8).expect("default concurrency is non-zero");
 
@@ -41,13 +42,15 @@ const DEFAULT_COLLECT_PREFETCH_CONCURRENCY: NonZeroUsize =
 #[derive(Clone, Debug)]
 pub struct QueryConfig {
     query_response_timeout: Duration,
+    query_cancel_request_timeout: Duration,
     collect_prefetch_concurrency: NonZeroUsize,
 }
 
 impl Default for QueryConfig {
     fn default() -> Self {
         Self {
             query_response_timeout: DEFAULT_QUERY_RESPONSE_TIMEOUT,
+            query_cancel_request_timeout: DEFAULT_QUERY_CANCEL_REQUEST_TIMEOUT,
             collect_prefetch_concurrency: DEFAULT_COLLECT_PREFETCH_CONCURRENCY,
         }
     }
@@ -229,6 +232,15 @@ impl QueryConfig {
         self
     }
 
+    /// Sets the client-side deadline for an explicit query cancellation request. Defaults to `30s`.
+    ///
+    /// The deadline covers abort request transport, response body reading, and bounded transport retries. It is
+    /// independent of [`Self::with_query_response_timeout`].
+    pub fn with_query_cancel_request_timeout(mut self, timeout: Duration) -> Self {
+        self.query_cancel_request_timeout = timeout;
+        self
+    }
+
     /// Sets the default number of partitions fetched concurrently during collection. Defaults to `8`.
     pub fn with_collect_prefetch_concurrency(mut self, concurrency: NonZeroUsize) -> Self {
         self.collect_prefetch_concurrency = concurrency;
@@ -324,6 +336,7 @@ impl From<SessionConfig> for InitialSessionConfig {
 #[derive(Debug)]
 pub(crate) struct QueryExecutionPolicy {
     query_response_timeout: Duration,
+    query_cancel_request_timeout: Duration,
     collect_prefetch_concurrency: NonZeroUsize,
 }
 
@@ -334,6 +347,9 @@ impl QueryExecutionPolicy {
             query_response_timeout: options
                 .query_response_timeout
                 .unwrap_or(self.query_response_timeout),
+            query_cancel_request_timeout: options
+                .query_cancel_request_timeout
+                .unwrap_or(self.query_cancel_request_timeout),
             collect_prefetch_concurrency: options
                 .collect_prefetch_concurrency
                 .unwrap_or(self.collect_prefetch_concurrency),
@@ -345,13 +361,15 @@ impl QueryExecutionPolicy {
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct QueryExecutionSettings {
     pub(crate) query_response_timeout: Duration,
+    pub(crate) query_cancel_request_timeout: Duration,
     pub(crate) collect_prefetch_concurrency: NonZeroUsize,
 }
 
 impl From<QueryConfig> for QueryExecutionPolicy {
     fn from(config: QueryConfig) -> Self {
         Self {
             query_response_timeout: config.query_response_timeout,
+            query_cancel_request_timeout: config.query_cancel_request_timeout,
             collect_prefetch_concurrency: config.collect_prefetch_concurrency,
         }
     }
@@ -498,6 +516,10 @@ mod tests {
         let policy: QueryExecutionPolicy = QueryConfig::default().into();
         let settings = policy.resolve_options(QueryOptions::default());
         assert_eq!(settings.query_response_timeout, Duration::from_secs(300));
+        assert_eq!(
+            settings.query_cancel_request_timeout,
+            Duration::from_secs(30)
+        );
     }
 
     #[test]
@@ -542,6 +564,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn query_options_with_query_cancel_request_timeout_overrides_only_cancel_timeout() {
+        let policy: QueryExecutionPolicy = QueryConfig::default()
+            .with_query_response_timeout(Duration::from_secs(120))
+            .with_collect_prefetch_concurrency(NonZeroUsize::new(4).unwrap())
+            .into();
+        let settings = policy.resolve_options(
+            QueryOptions::default().with_query_cancel_request_timeout(Duration::from_secs(7)),
+        );
+
+        assert_eq!(settings.query_response_timeout, Duration::from_secs(120));
+        assert_eq!(
+            settings.query_cancel_request_timeout,
+            Duration::from_secs(7)
+        );
+        assert_eq!(
+            settings.collect_prefetch_concurrency,
+            NonZeroUsize::new(4).unwrap()
+        );
+    }
+
     #[test]
     fn query_options_with_collect_prefetch_concurrency_overrides_only_concurrency() {
         let policy: QueryExecutionPolicy = QueryConfig::default()

src/error.rs

@@ -23,10 +23,12 @@ pub use schema::{
 };
 
 pub(crate) use parse::RowsetParseError;
-pub(crate) use query_scoped::{QueryScopedError, QueryScopedRepr, QueryScopedResult};
+pub(crate) use query_scoped::{
+    QueryScopedError, QueryScopedRepr, QueryScopedResult, with_optional_query_id,
+};
 pub(crate) use repr::{
-    AuthError, ConfigError, InternalError, NetworkError, ProtocolError, ServerError,
-    SessionExpiredError, TimeoutError,
+    AuthError, CancelledError, ConfigError, InternalError, NetworkError, ProtocolError,
+    ServerError, SessionExpiredError, TimeoutError,
 };
 
 const VALUE_PREVIEW_MAX_CHARS: usize = 128;
@@ -44,6 +46,7 @@ const JSON_BODY_PREVIEW_MAX_BYTES: usize = 1024;
 ///
 /// - Snowflake-provided fields: [`snowflake_code`](Error::snowflake_code),
 ///   [`snowflake_message`](Error::snowflake_message), [`query_id`](Error::query_id)
+/// - Cancellation failures: [`is_cancelled`](Error::is_cancelled)
 /// - Decode failures: [`as_schema_error`](Error::as_schema_error) for schema validation failures,
 ///   [`as_cell_decode_error`](Error::as_cell_decode_error) for cell conversion failures,
 ///   [`as_custom_plan_error`](Error::as_custom_plan_error) for plan-time failures raised by a hand-written decoder,
@@ -119,6 +122,8 @@ pub enum ErrorKind {
     Network,
     /// Snowflake rejected the request and returned a server-side error message.
     Server,
+    /// A user-requested cancellation prevented submission or terminated the query.
+    Cancelled,
     /// The current session token is no longer valid.
     SessionExpired,
     /// The connector timed out while waiting for a response.
@@ -157,6 +162,7 @@ impl Error {
             Repr::Auth(_) => ErrorKind::Auth,
             Repr::Network { .. } => ErrorKind::Network,
             Repr::Server(_) => ErrorKind::Server,
+            Repr::Cancelled(_) => ErrorKind::Cancelled,
             Repr::SessionExpired(_) => ErrorKind::SessionExpired,
             Repr::Timeout { .. } => ErrorKind::Timeout,
             Repr::Protocol { .. } => ErrorKind::Protocol,
@@ -176,6 +182,7 @@ impl Error {
             Repr::Auth(AuthError::LoginRejected { message }) => message.as_deref(),
             Repr::SessionExpired(SessionExpiredError { message, .. }) => message.as_deref(),
             Repr::Server(ServerError { message, .. }) => message.as_deref(),
+            Repr::Cancelled(CancelledError { message, .. }) => message.as_deref(),
             _ => None,
         }
     }
@@ -185,6 +192,7 @@ impl Error {
         match &*self.repr {
             Repr::SessionExpired(SessionExpiredError { code, .. }) => code.as_deref(),
             Repr::Server(ServerError { code, .. }) => code.as_deref(),
+            Repr::Cancelled(CancelledError { code, .. }) => code.as_deref(),
             _ => None,
         }
     }
@@ -204,6 +212,7 @@ impl Error {
             | Repr::Protocol { query_id, .. }
             | Repr::Internal { query_id, .. } => query_id.as_deref(),
             Repr::Server(ServerError { query_id, .. }) => query_id.as_deref(),
+            Repr::Cancelled(CancelledError { query_id, .. }) => query_id.as_deref(),
             Repr::SessionExpired(SessionExpiredError { query_id, .. }) => query_id.as_deref(),
             _ => None,
         }
@@ -221,6 +230,11 @@ impl Error {
         matches!(&*self.repr, Repr::Server(_))
     }
 
+    /// Reports whether this error is a cancellation, i.e. [`ErrorKind::Cancelled`].
+    pub fn is_cancelled(&self) -> bool {
+        matches!(&*self.repr, Repr::Cancelled(_))
+    }
+
     pub fn as_cell_decode_error(&self) -> Option<&CellDecodeError> {
         match &*self.repr {
             Repr::CellDecode(error) => Some(error),
@@ -328,6 +342,12 @@ impl From<ServerError> for Error {
     }
 }
 
+impl From<CancelledError> for Error {
+    fn from(error: CancelledError) -> Self {
+        Self::new(Repr::Cancelled(error))
+    }
+}
+
 impl From<SessionExpiredError> for Error {
     fn from(error: SessionExpiredError) -> Self {
         Self::new(Repr::SessionExpired(error))
@@ -500,6 +520,20 @@ impl ServerError {
     }
 }
 
+impl CancelledError {
+    pub(crate) fn new(
+        code: Option<String>,
+        message: Option<String>,
+        query_id: Option<Arc<str>>,
+    ) -> Self {
+        Self {
+            code: code.map(String::into_boxed_str),
+            message: message.map(String::into_boxed_str),
+            query_id,
+        }
+    }
+}
+
 impl SessionExpiredError {
     pub(crate) fn new(
         code: Option<String>,
@@ -527,6 +561,10 @@ impl TimeoutError {
         Self::Query
     }
 
+    pub(crate) fn query_cancel() -> Self {
+        Self::QueryCancel
+    }
+
     #[cfg(feature = "external-browser-sso")]
     pub(crate) fn browser_callback() -> Self {
         Self::BrowserCallback
@@ -745,6 +783,37 @@ mod tests {
         assert_eq!(server_err.to_string(), "Snowflake server error 390100");
     }
 
+    #[test]
+    fn cancelled_error_exposes_structured_details() {
+        let err: Error = CancelledError::new(
+            Some("000604".to_string()),
+            Some("SQL execution canceled".to_string()),
+            Some(Arc::from("query-id")),
+        )
+        .into();
+
+        assert_eq!(err.kind(), ErrorKind::Cancelled);
+        assert!(err.is_cancelled());
+        assert_eq!(err.snowflake_code(), Some("000604"));
+        assert_eq!(err.snowflake_message(), Some("SQL execution canceled"));
+        assert_eq!(err.query_id(), Some("query-id"));
+        assert_eq!(
+            err.to_string(),
+            "query cancelled: SQL execution canceled (query id: query-id)"
+        );
+    }
+
+    #[test]
+    fn cancelled_error_without_server_fields_has_stable_display() {
+        let err: Error = CancelledError::new(None, None, None).into();
+        assert_eq!(err.kind(), ErrorKind::Cancelled);
+        assert!(err.is_cancelled());
+        assert_eq!(err.to_string(), "query cancelled");
+        assert_eq!(err.snowflake_code(), None);
+        assert_eq!(err.snowflake_message(), None);
+        assert_eq!(err.query_id(), None);
+    }
+
     #[test]
     fn session_expired_preserves_snowflake_fields() {
         let err: Error = SessionExpiredError::new(

src/error/display.rs

@@ -6,8 +6,8 @@ use std::{
 use super::{
     Error,
     repr::{
-        AuthError, ConfigError, InternalError, NetworkError, ProtocolError, Repr, ServerError,
-        TimeoutError,
+        AuthError, CancelledError, ConfigError, InternalError, NetworkError, ProtocolError, Repr,
+        ServerError, TimeoutError,
     },
 };
 
@@ -79,6 +79,18 @@ impl Display for Error {
                 }
                 Ok(())
             }
+            Repr::Cancelled(CancelledError {
+                message, query_id, ..
+            }) => {
+                f.write_str("query cancelled")?;
+                if let Some(message) = message {
+                    write!(f, ": {message}")?;
+                }
+                if let Some(query_id) = query_id {
+                    write!(f, " (query id: {query_id})")?;
+                }
+                Ok(())
+            }
             Repr::SessionExpired(_) => f.write_str("session expired"),
             Repr::Timeout {
                 error: TimeoutError::Request(_),
@@ -88,6 +100,10 @@ impl Display for Error {
                 error: TimeoutError::Query,
                 ..
             } => f.write_str("timed out waiting for query response"),
+            Repr::Timeout {
+                error: TimeoutError::QueryCancel,
+                ..
+            } => f.write_str("timed out waiting for query cancellation response"),
             #[cfg(feature = "external-browser-sso")]
             Repr::Timeout {
                 error: TimeoutError::BrowserCallback,