From 5110e7d54b92522907703c1d2978f684213f73a6 Mon Sep 17 00:00:00 2001 From: SorenHolstHansen Date: Thu, 20 Mar 2025 12:39:05 +0100 Subject: [PATCH 1/4] chore: add more fields to responses and JsonResult --- snowflake-api/src/lib.rs | 11 +++++++++-- snowflake-api/src/responses.rs | 19 ++++++++++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index 6ce9e5d..6d76df2 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -28,7 +28,7 @@ use reqwest_middleware::ClientWithMiddleware; use serde_json::Value; use thiserror::Error; -use responses::{ExecResponse, ExecRestResponse, ProcessedRestResponse}; +use responses::{ExecResponse, ExecRestResponse, ProcessedRestResponse, QueryContext}; use session::{AuthError, Session}; use crate::connection::QueryType; @@ -116,6 +116,9 @@ pub struct JsonResult { pub value: serde_json::Value, /// Field ordering matches the array ordering pub schema: Vec, + pub query_id: String, + pub send_result_time: usize, + pub query_context: QueryContext, } impl Display for JsonResult { @@ -423,7 +426,8 @@ impl SnowflakeApi { match resp { ExecResponse::Query(_) => Err(SnowflakeApiError::UnexpectedResponse), ExecResponse::PutGet(pg) => { - let res = into_resp_type!(&pg, RawQueryResult::Empty(EmptyJsonResult { schema: None })); + let res = + into_resp_type!(&pg, RawQueryResult::Empty(EmptyJsonResult { schema: None })); put::put(pg).await?; Ok(res) } @@ -521,6 +525,9 @@ impl SnowflakeApi { let value = serde_json::to_value(values).unwrap(); RawQueryResult::Json(JsonResult { value, + query_id: sync_data.query_id, + send_result_time: sync_data.send_result_time, + query_context: sync_data.query_context, schema: sync_data .rowtype .unwrap() diff --git a/snowflake-api/src/responses.rs b/snowflake-api/src/responses.rs index 969d385..2647e3b 100644 --- a/snowflake-api/src/responses.rs +++ b/snowflake-api/src/responses.rs @@ -177,6 +177,21 @@ pub struct AsyncQueryExecResponseData { pub query_aborts_after_secs: i64, } +#[derive(Deserialize, Debug, Clone)] +#[serde(rename_all = "camelCase")] +pub struct QueryContextEntry { + id: isize, + timestamp: usize, + // priority: usize, + // context: String +} + +#[derive(Deserialize, Debug, Clone)] +#[serde(rename_all = "camelCase")] +pub struct QueryContext { + entries: Vec, +} + #[derive(Deserialize, Debug, Clone)] #[serde(rename_all = "camelCase")] pub struct SyncQueryExecResponseData { @@ -212,7 +227,9 @@ pub struct SyncQueryExecResponseData { // multi-statement response, comma-separated pub result_ids: Option, // `progressDesc`, and `queryAbortAfterSecs` are not used but exist in .NET - // `sendResultTime`, `queryResultFormat`, `queryContext` also exist + // `queryResultFormat` also exist + pub send_result_time: usize, + pub query_context: QueryContext, } #[derive(Deserialize, Debug, Clone)] From 5894487a44fce581cbd34f1cbdcc415b01d47b2b Mon Sep 17 00:00:00 2001 From: SorenHolstHansen Date: Thu, 20 Mar 2025 12:51:33 +0100 Subject: [PATCH 2/4] chore: make fields public --- snowflake-api/src/responses.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/snowflake-api/src/responses.rs b/snowflake-api/src/responses.rs index 2647e3b..2b92004 100644 --- a/snowflake-api/src/responses.rs +++ b/snowflake-api/src/responses.rs @@ -180,16 +180,16 @@ pub struct AsyncQueryExecResponseData { #[derive(Deserialize, Debug, Clone)] #[serde(rename_all = "camelCase")] pub struct QueryContextEntry { - id: isize, - timestamp: usize, - // priority: usize, - // context: String + pub id: isize, + pub timestamp: usize, + // pub priority: usize, + // pub context: String } #[derive(Deserialize, Debug, Clone)] #[serde(rename_all = "camelCase")] pub struct QueryContext { - entries: Vec, + pub entries: Vec, } #[derive(Deserialize, Debug, Clone)] From bf961f296d1090243e2d2f6c3f7d6a7c7b5d6e5a Mon Sep 17 00:00:00 2001 From: SorenHolstHansen Date: Thu, 20 Mar 2025 13:16:14 +0100 Subject: [PATCH 3/4] chore: add the fields to arrow responses --- snowflake-api/src/lib.rs | 39 +++++++++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index 6d76df2..c3606a3 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -127,6 +127,14 @@ impl Display for JsonResult { } } +#[derive(Debug)] +pub struct BytesResult { + pub chunks: Vec, + pub query_id: String, + pub send_result_time: usize, + pub query_context: QueryContext, +} + /// Based on the [`ExecResponseRowType`] #[derive(Debug)] pub struct FieldSchema { @@ -150,12 +158,20 @@ impl From for FieldSchema { } } +#[derive(Debug)] +pub struct ArrowResult { + pub batches: Vec, + pub query_id: String, + pub send_result_time: usize, + pub query_context: QueryContext, +} + /// Container for query result. /// Arrow is returned by-default for all SELECT statements, /// unless there is session configuration issue or it's a different statement type. #[derive(Debug)] pub enum QueryResult { - Arrow(Vec), + Arrow(ArrowResult), Json(JsonResult), Empty(EmptyJsonResult), } @@ -165,7 +181,7 @@ pub enum QueryResult { pub enum RawQueryResult { /// Arrow IPC chunks /// see: - Bytes(Vec), + Bytes(BytesResult), /// Json payload is deserialized, /// as it's already a part of REST response Json(JsonResult), @@ -175,9 +191,15 @@ pub enum RawQueryResult { impl RawQueryResult { pub fn deserialize_arrow(self) -> Result { match self { - RawQueryResult::Bytes(bytes) => { - Self::flat_bytes_to_batches(bytes).map(QueryResult::Arrow) - } + RawQueryResult::Bytes(bytes_result) => Self::flat_bytes_to_batches(bytes_result.chunks) + .map(|batches| { + QueryResult::Arrow(ArrowResult { + batches, + query_id: bytes_result.query_id, + send_result_time: bytes_result.send_result_time, + query_context: bytes_result.query_context, + }) + }), RawQueryResult::Json(j) => Ok(QueryResult::Json(j)), RawQueryResult::Empty(e) => Ok(QueryResult::Empty(e)), } @@ -551,7 +573,12 @@ impl SnowflakeApi { chunks.push(bytes); } - RawQueryResult::Bytes(chunks) + RawQueryResult::Bytes(BytesResult { + chunks, + query_id: sync_data.query_id, + send_result_time: sync_data.send_result_time, + query_context: sync_data.query_context, + }) } else { return Err(SnowflakeApiError::BrokenResponse); }; From f1d0504d485d6d9241482c4bb1204b9b01388eea Mon Sep 17 00:00:00 2001 From: SorenHolstHansen Date: Thu, 20 Mar 2025 13:22:38 +0100 Subject: [PATCH 4/4] chore: add the fields to the empty and arrow types as well --- snowflake-api/src/lib.rs | 21 ++++++++++++++++++--- snowflake-api/src/responses.rs | 3 +++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index c3606a3..c46f2ba 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -106,6 +106,9 @@ pub enum SnowflakeApiError { #[derive(Debug)] pub struct EmptyJsonResult { pub schema: Option>, + pub query_id: String, + pub send_result_time: usize, + pub query_context: QueryContext, } /// Even if Arrow is specified as a return type non-select queries @@ -448,8 +451,15 @@ impl SnowflakeApi { match resp { ExecResponse::Query(_) => Err(SnowflakeApiError::UnexpectedResponse), ExecResponse::PutGet(pg) => { - let res = - into_resp_type!(&pg, RawQueryResult::Empty(EmptyJsonResult { schema: None })); + let res = into_resp_type!( + &pg, + RawQueryResult::Empty(EmptyJsonResult { + schema: None, + query_id: pg.data.query_id.clone(), + send_result_time: pg.data.send_result_time, + query_context: pg.data.query_context.clone() + }) + ); put::put(pg).await?; Ok(res) } @@ -518,7 +528,12 @@ impl SnowflakeApi { } else { None }; - RawQueryResult::Empty(EmptyJsonResult { schema }) + RawQueryResult::Empty(EmptyJsonResult { + schema, + query_id: sync_data.query_id, + send_result_time: sync_data.send_result_time, + query_context: sync_data.query_context, + }) } else if let Some(value) = sync_data.rowset { log::debug!("Got JSON response"); let mut values: Vec = serde_json::from_value(value).unwrap(); diff --git a/snowflake-api/src/responses.rs b/snowflake-api/src/responses.rs index 2b92004..69acdb9 100644 --- a/snowflake-api/src/responses.rs +++ b/snowflake-api/src/responses.rs @@ -300,6 +300,9 @@ pub struct PutGetResponseData { #[serde(default)] pub parameters: Vec, pub statement_type_id: Option, + pub query_id: String, + pub send_result_time: usize, + pub query_context: QueryContext, } #[derive(Deserialize, Debug)]