snowflake rest retries support (#24)

* replace oneshot by JoinHandle

* switch to different execution model with results waiting until task joined

* fix abort

* simplify executor interface

Author: YaroslavLitvinov

crates/api-snowflake-rest/src/server/error.rs

@@ -6,7 +6,7 @@ use datafusion::arrow::error::ArrowError;
 use error_stack::ErrorChainExt;
 use error_stack::ErrorExt;
 use error_stack_trace;
-use executor::QueryRecordId;
+use executor::QueryId;
 use executor::error::OperationOn;
 use executor::error_code::ErrorCode;
 use executor::snowflake_error::Entity;
@@ -132,11 +132,11 @@ impl Error {
         InvalidAuthDataSnafu.build()
     }
 
-    pub fn query_id(&self) -> QueryRecordId {
+    pub fn query_id(&self) -> QueryId {
         if let Self::Execution { source, .. } = self {
             source.query_id()
         } else {
-            QueryRecordId::default()
+            QueryId::default()
         }
     }
 
@@ -236,8 +236,7 @@ impl Error {
         tracing::Span::current()
             .record("error_code", error_code.to_string())
             .record("sql_state", sql_state.to_string())
-            .record("query_id", self.query_id().as_i64())
-            .record("query_uuid", self.query_id().as_uuid().to_string())
+            .record("query_id", self.query_id().to_string())
             .record("display_error", &display_error)
             .record("debug_error", self.debug_error_message())
             .record("error_stack_trace", self.output_msg())
@@ -260,7 +259,7 @@ impl Error {
                 returned: None,
                 query_result_format: None,
                 // Query uuid is returned to the user
-                query_id: Some(self.query_id().as_uuid().to_string()),
+                query_id: Some(self.query_id().to_string()),
             }),
             code: Some(error_code.to_string()),
         });

crates/api-snowflake-rest/src/server/handlers.rs

@@ -34,7 +34,7 @@ pub async fn login(
     name = "api_snowflake_rest::query",
     level = "debug",
     skip(state),
-    fields(query_id, query_uuid),
+    fields(query_id),
     err,
     ret(level = tracing::Level::TRACE),
 )]
@@ -65,9 +65,10 @@ pub async fn abort(
         request_id,
     }): Json<AbortRequestBody>,
 ) -> Result<Json<serde_json::value::Value>> {
-    state
+    let query_id = state
         .execution_svc
-        .abort_query(RunningQueryId::ByRequestId(request_id, sql_text))?;
+        .locate_query_id(RunningQueryId::ByRequestId(request_id, sql_text))?;
+    state.execution_svc.abort(query_id)?;
     Ok(Json(serde_json::value::Value::Null))
 }
 

crates/api-snowflake-rest/src/server/helpers.rs

@@ -74,7 +74,7 @@ fn records_to_json_string(recs: &[RecordBatch]) -> std::result::Result<String, E
 )]
 pub fn handle_query_ok_result(
     sql_text: &str,
-    query_uuid: Uuid,
+    query_id: Uuid,
     query_result: QueryResult,
     ser_fmt: DataSerializationFormat,
 ) -> Result<JsonResponse> {
@@ -115,7 +115,7 @@ pub fn handle_query_ok_result(
             row_set_base_64,
             total: Some(total_rows),
             returned: Some(returned_rows),
-            query_id: Some(query_uuid.to_string()),
+            query_id: Some(query_id.to_string()),
             error_code: None,
             sql_state: Some(SqlState::Success.to_string()),
         }),

crates/api-snowflake-rest/src/server/logic.rs

@@ -8,6 +8,7 @@ use crate::server::error::{
 };
 use crate::server::helpers::handle_query_ok_result;
 use api_snowflake_rest_sessions::helpers::{create_jwt, ensure_jwt_secret_is_valid, jwt_claims};
+use executor::RunningQueryId;
 use executor::models::QueryContext;
 use snafu::{OptionExt, ResultExt};
 use time::Duration;
@@ -87,11 +88,28 @@ pub async fn handle_query_request(
         return api_snowflake_rest_error::NotImplementedSnafu.fail();
     }
 
-    let query_uuid = query_context.query_id.as_uuid();
-    let result = state
-        .execution_svc
-        .query(session_id, &sql_text, query_context)
-        .await?;
+    // find running query by request_id
+    let session = state.execution_svc.get_session(session_id).await?;
+    let query_id_res = session
+        .running_queries
+        .locate_query_id(RunningQueryId::ByRequestId(
+            query.request_id,
+            sql_text.clone(),
+        ));
 
-    handle_query_ok_result(&sql_text, query_uuid, result, serialization_format)
+    let (result, query_id) = if query.retry_count.unwrap_or_default() > 0
+        && let Ok(query_id) = query_id_res
+    {
+        let result = state.execution_svc.wait(query_id).await?;
+        (result, query_id)
+    } else {
+        let query_id = query_context.query_id;
+        let result = state
+            .execution_svc
+            .query(session_id, &sql_text, query_context)
+            .await?;
+        (result, query_id)
+    };
+
+    handle_query_ok_result(&sql_text, query_id, result, serialization_format)
 }

crates/executor/src/datafusion/rewriters/session_context.rs

@@ -1,6 +1,6 @@
 use crate::models::QueryContext;
-use crate::query_types::QueryRecordId;
-use crate::running_queries::{RunningQueries, RunningQueryId};
+use crate::query_types::QueryId;
+use crate::running_queries::RunningQueries;
 use datafusion::arrow::array::{ListArray, ListBuilder, StringBuilder};
 use datafusion::logical_expr::{Expr, LogicalPlan};
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
@@ -19,7 +19,7 @@ pub struct SessionContextExprRewriter {
     pub session_id: String,
     pub version: String,
     pub query_context: QueryContext,
-    pub recent_queries: Arc<RwLock<VecDeque<QueryRecordId>>>,
+    pub recent_queries: Arc<RwLock<VecDeque<QueryId>>>,
     pub running_queries: Arc<dyn RunningQueries>,
 }
 
@@ -64,24 +64,17 @@ impl SessionContextExprRewriter {
 
     #[allow(clippy::needless_pass_by_value)]
     pub fn cancel_query(&self, query_id: String) -> Result<ScalarValue> {
-        let Ok(query_uuid) = Uuid::from_str(&query_id) else {
+        let Ok(query_id) = Uuid::from_str(&query_id) else {
             return Ok(utf8_val("Invalid UUID."));
         };
 
-        let query_id = QueryRecordId::from(query_uuid);
-        if self
-            .running_queries
-            .abort(RunningQueryId::ByQueryId(query_id))
-            .is_err()
-        {
+        let query_id = QueryId::from(query_id);
+        if self.running_queries.abort(query_id).is_err() {
             Ok(utf8_val(
                 "Identified SQL statement is not currently executing.",
             ))
         } else {
-            Ok(utf8_val(format!(
-                "query [{}] terminated.",
-                query_id.as_uuid()
-            )))
+            Ok(utf8_val(format!("query [{query_id}] terminated.")))
         }
     }
 }

crates/executor/src/error.rs

@@ -1,5 +1,5 @@
 use super::snowflake_error::SnowflakeError;
-use crate::query_types::{QueryRecordId, QueryStatus};
+use crate::query_types::{QueryId, QueryStatus};
 use catalog::error::Error as CatalogError;
 use datafusion_common::DataFusionError;
 use error_stack_trace;
@@ -544,18 +544,18 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("{}: Query execution error: {source}", query_id.as_uuid()))]
+    #[snafu(display("{query_id}: Query execution error: {source}"))]
     QueryExecution {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(source(from(Error, Box::new)))]
         source: Box<Error>,
         #[snafu(implicit)]
         location: Location,
     },
 
-    #[snafu(display("Query {} isn't running", query_id.as_uuid()))]
+    #[snafu(display("Query {query_id} isn't running"))]
     QueryIsntRunning {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(implicit)]
         location: Location,
     },
@@ -568,48 +568,32 @@ pub enum Error {
     },
 
     // When user tried to get result before query finished
-    #[snafu(display("Query {} is running", query_id.as_uuid()))]
+    #[snafu(display("Query {query_id} is running"))]
     QueryIsRunning {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(implicit)]
         location: Location,
     },
 
-    #[snafu(display("Query {} cancelled", query_id.as_uuid()))]
+    #[snafu(display("Query {query_id} cancelled"))]
     QueryCancelled {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(implicit)]
         location: Location,
     },
 
-    #[snafu(display("Query [{}] result sending error", query_id.as_uuid()))]
-    QueryResultSend {
-        query_id: QueryRecordId,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
-    #[snafu(display("Query [{}] result recv error: {error}", query_id.as_uuid()))]
-    QueryResultRecv {
-        query_id: QueryRecordId,
-        #[snafu(source)]
-        error: tokio::sync::oneshot::error::RecvError,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
-    #[snafu(display("Query [{}] result notify error: {error}", query_id.as_uuid()))]
+    #[snafu(display("Query {query_id} result notify error: {error}"))]
     QueryStatusRecv {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(source)]
         error: tokio::sync::watch::error::RecvError,
         #[snafu(implicit)]
         location: Location,
     },
 
-    #[snafu(display("Query [{}] status notify error: {error}", query_id.as_uuid()))]
+    #[snafu(display("Query {query_id} status notify error: {error}"))]
     NotifyQueryStatus {
-        query_id: QueryRecordId,
+        query_id: QueryId,
         #[snafu(source)]
         error: tokio::sync::watch::error::SendError<QueryStatus>,
         #[snafu(implicit)]
@@ -624,14 +608,38 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Failed to get async result for query [{query_id}]: {error}"))]
+    AsyncResultTaskJoin {
+        #[snafu(source)]
+        error: tokio::task::JoinError,
+        query_id: QueryId,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display("Failed to join query subtask: {error}"))]
+    QuerySubtaskJoin {
+        #[snafu(source)]
+        error: tokio::task::JoinError,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display("Missing results handle in running query [{query_id:?}]"))]
+    NoJoinHandle {
+        query_id: QueryId,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 impl Error {
-    pub fn query_id(&self) -> QueryRecordId {
+    pub fn query_id(&self) -> QueryId {
         if let Self::QueryExecution { query_id, .. } = self {
             *query_id
         } else {
-            QueryRecordId::default()
+            QueryId::default()
         }
     }
     #[must_use]

crates/executor/src/lib.rs

@@ -17,7 +17,7 @@ pub mod utils;
 pub mod tests;
 
 pub use error::{Error, Result};
-pub use query_types::{QueryRecordId, QueryStatus};
+pub use query_types::{QueryId, QueryStatus};
 pub use running_queries::RunningQueryId;
 pub use snowflake_error::SnowflakeError;
 

crates/executor/src/models.rs

@@ -1,25 +1,19 @@
-use crate::query_types::{QueryRecordId, QueryStatus};
+use crate::query_types::QueryId;
 use datafusion::arrow::array::RecordBatch;
 use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit};
 use datafusion_common::arrow::datatypes::Schema;
 use functions::to_snowflake_datatype;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::sync::Arc;
-use tokio::sync::oneshot;
 use uuid::Uuid;
 
-pub struct AsyncQueryHandle {
-    pub query_id: QueryRecordId,
-    pub rx: oneshot::Receiver<QueryResultStatus>,
-}
-
 #[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 pub struct QueryContext {
     pub database: Option<String>,
     pub schema: Option<String>,
     pub worksheet_id: Option<i64>,
-    pub query_id: QueryRecordId,
+    pub query_id: QueryId,
     pub request_id: Option<Uuid>,
     pub ip_address: Option<String>,
 }
@@ -35,14 +29,14 @@ impl QueryContext {
             database,
             schema,
             worksheet_id,
-            query_id: QueryRecordId::default(),
+            query_id: QueryId::default(),
             request_id: None,
             ip_address: None,
         }
     }
 
     #[must_use]
-    pub const fn with_query_id(mut self, new_id: QueryRecordId) -> Self {
+    pub const fn with_query_id(mut self, new_id: QueryId) -> Self {
         self.query_id = new_id;
         self
     }
@@ -80,12 +74,6 @@ impl QueryResult {
     }
 }
 
-#[derive(Debug)]
-pub struct QueryResultStatus {
-    pub query_result: std::result::Result<QueryResult, crate::Error>,
-    pub status: QueryStatus,
-}
-
 // TODO: We should not have serde dependency here
 // Instead it should be in api-snowflake-rest
 #[derive(Debug, Serialize, Deserialize, Clone)]

crates/executor/src/query.rs

@@ -262,8 +262,7 @@ impl UserQuery {
         skip(self),
         fields(
             statement,
-            query_id = self.query_context.query_id.as_i64(),
-            query_uuid = self.query_context.query_id.as_uuid().to_string(),
+            query_id = self.query_context.query_id.to_string(),
         ),
         err
     )]