Add Statestore Query to ExecutionService

Cargo.toml

@@ -92,7 +92,7 @@ url = "2.5"
 utoipa = { version = "5.3.1", features = ["uuid", "chrono"] }
 utoipa-axum = { version = "0.2.0" }
 utoipa-swagger-ui = { version = "9", features = ["axum"] }
-uuid = { version = "1.10.0", features = ["v4", "serde"] }
+uuid = { version = "1.10.0", features = ["v4", "v7", "serde"] }
 validator = { version = "0.20.0", features = ["derive"] }
 mockall = "0.13.1"
 insta = { version = "1.44.1", features = ["json", "filters", "redactions"] }

crates/api-snowflake-rest/Cargo.toml

@@ -7,6 +7,7 @@ license-file.workspace = true
 [features]
 default = []
 retry-disable = []
+traces-test-log = []
 
 [dependencies]
 api-snowflake-rest-sessions = { path = "../api-snowflake-rest-sessions" }

crates/api-snowflake-rest/src/tests/create_test_server.rs

@@ -12,7 +12,8 @@ use std::sync::{Arc, Condvar, Mutex};
 use std::thread;
 use std::time::Duration;
 use tokio::runtime::Builder;
-use tracing_subscriber::fmt::format::FmtSpan;
+#[cfg(feature = "traces-test-log")]
+use tracing_subscriber::{fmt, fmt::format::FmtSpan};
 
 static INIT: std::sync::Once = std::sync::Once::new();
 
@@ -140,10 +141,13 @@ fn setup_tracing() {
                             .with_targets(targets_with_level(&DISABLED_TARGETS, LevelFilter::OFF))
                             .with_default(LevelFilter::TRACE),
                     ),
-            )
+            );
+
+        #[cfg(feature = "traces-test-log")]
+        let registry = registry
             // Logs filtering
             .with(
-                tracing_subscriber::fmt::layer()
+                fmt::layer()
                     .with_writer(
                         std::fs::OpenOptions::new()
                             .create(true)

crates/build-info/build.rs

@@ -45,10 +45,11 @@ fn main() {
     println!("cargo:rustc-env=BUILD_TIMESTAMP={build_timestamp}");
 
     // Rerun build script if git HEAD changes
-    println!("cargo:rerun-if-changed=.git/HEAD");
+    // Should point to the root of the repository
+    println!("cargo:rerun-if-changed=../../.git/HEAD");
     // Also rerun if the current branch ref changes
     if let Some(branch_ref) = run_git_command(&["symbolic-ref", "HEAD"]) {
-        let ref_path = format!(".git/{branch_ref}");
+        let ref_path = format!("../../.git/{branch_ref}");
         println!("cargo:rerun-if-changed={ref_path}");
     }
 }

crates/embucket-lambda/Cargo.toml

@@ -45,6 +45,7 @@ retry-disable = ["api-snowflake-rest/retry-disable"]
 streaming = []
 rest-catalog = ["executor/rest-catalog"]
 dedicated-executor = ["executor/dedicated-executor"]
+state-store-query = ["executor/state-store-query"]
 
 [package.metadata.lambda]
 # Default binary to deploy

crates/embucketd/Cargo.toml

@@ -52,3 +52,4 @@ retry-disable = ["api-snowflake-rest/retry-disable"]
 rest-catalog = ["executor/rest-catalog"]
 dedicated-executor = ["executor/dedicated-executor"]
 state-store = ["executor/state-store"]
+state-store-query = ["executor/state-store-query"]

crates/executor/Cargo.toml

@@ -5,9 +5,11 @@ edition = "2024"
 license-file.workspace = true
 
 [features]
+default = []
 rest-catalog = ["catalog/rest-catalog"]
 dedicated-executor = []
 state-store = []
+state-store-query = ["state-store"]
 
 [dependencies]
 catalog-metastore = { path = "../catalog-metastore" }

crates/executor/src/error.rs

@@ -1,5 +1,5 @@
 use super::snowflake_error::SnowflakeError;
-use crate::query_types::{QueryId, QueryStatus};
+use crate::query_types::{ExecutionStatus, QueryId};
 use catalog::error::Error as CatalogError;
 use datafusion_common::DataFusionError;
 use error_stack_trace;
@@ -592,7 +592,7 @@ pub enum Error {
     },
 
     #[snafu(display("Query {query_id} result notify error: {error}"))]
-    QueryStatusRecv {
+    ExecutionStatusRecv {
         query_id: QueryId,
         #[snafu(source)]
         error: tokio::sync::watch::error::RecvError,
@@ -601,10 +601,10 @@ pub enum Error {
     },
 
     #[snafu(display("Query {query_id} status notify error: {error}"))]
-    NotifyQueryStatus {
+    NotifyExecutionStatus {
         query_id: QueryId,
         #[snafu(source)]
-        error: tokio::sync::watch::error::SendError<QueryStatus>,
+        error: tokio::sync::watch::error::SendError<ExecutionStatus>,
         #[snafu(implicit)]
         location: Location,
     },

crates/executor/src/error_code.rs

@@ -6,8 +6,10 @@ use std::fmt::Display;
 // For reference: https://github.com/snowflakedb/snowflake-cli/blob/main/src/snowflake/cli/api/errno.py
 // Some of our error codes may be mapped to Snowflake error codes
 
+// Do not set values for error codes, they are assigned in Display trait
 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
 pub enum ErrorCode {
+    None,
     Db,
     Metastore,
     #[cfg(feature = "state-store")]
@@ -35,13 +37,19 @@ pub enum ErrorCode {
     EntityNotFound(Entity, OperationOn),
     Other,
     UnsupportedFeature,
+    Timeout,
+    Cancelled,
+    LimitExceeded,
+    QueryTask,
 }
 
 impl Display for ErrorCode {
     #[allow(clippy::unnested_or_patterns)]
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let code = match self {
             Self::UnsupportedFeature => 2,
+            Self::Timeout => 630,
+            Self::Cancelled => 684,
             Self::HistoricalQueryError => 1001,
             Self::DataFusionSqlParse => 1003,
             Self::DataFusionSql => 2003,

crates/executor/src/lib.rs

@@ -6,6 +6,7 @@ pub mod error;
 pub mod error_code;
 pub mod models;
 pub mod query;
+pub mod query_task_result;
 pub mod query_types;
 pub mod running_queries;
 pub mod service;
@@ -18,7 +19,7 @@ pub mod utils;
 pub mod tests;
 
 pub use error::{Error, Result};
-pub use query_types::{QueryId, QueryStatus};
+pub use query_types::{ExecutionStatus, QueryId};
 pub use running_queries::RunningQueryId;
 pub use snowflake_error::SnowflakeError;
 

crates/executor/src/models.rs

@@ -8,7 +8,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use uuid::Uuid;
 
-#[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 pub struct QueryContext {
     pub database: Option<String>,
     pub schema: Option<String>,
@@ -18,6 +18,21 @@ pub struct QueryContext {
     pub ip_address: Option<String>,
 }
 
+// Add own Default implementation to avoid getting default (zeroed) Uuid.
+// This compromise is against rules, since this default is not deterministic.
+impl Default for QueryContext {
+    fn default() -> Self {
+        Self {
+            database: None,
+            schema: None,
+            worksheet_id: None,
+            query_id: Uuid::now_v7(),
+            request_id: None,
+            ip_address: None,
+        }
+    }
+}
+
 impl QueryContext {
     #[must_use]
     pub fn new(
@@ -29,9 +44,7 @@ impl QueryContext {
             database,
             schema,
             worksheet_id,
-            query_id: QueryId::default(),
-            request_id: None,
-            ip_address: None,
+            ..Default::default()
         }
     }