Files
windmill/backend/windmill-worker/src/snowflake_executor.rs
T
Ruben FiszelandClaude Opus 4.6 a6d4390790 feat: workflow-as-code (WAC) v2 (#8172)
* feat: workflow-as-code v2 with @task decorator API

Replace ctx.step("name", "script") API with @task decorators where
functions are called directly. Users no longer need to pass WorkflowCtx
or use string-based step names/script paths.

Python: @task decorator with contextvars-based implicit context
TypeScript: task() wrapper with module-level context variable
Parsers: detect @task function calls instead of ctx.step() calls
Worker: updated wrappers to set implicit context

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: WAC v2 checkpoint/replay with _executing_key child dispatch

- Rust-side orchestration: parent dispatches child jobs, suspends, resumes on completion
- _executing_key in checkpoint tells child which step to execute directly
- task() throws StepSuspend(mode="step_complete") after executing target step
- result_processor handles child completion and updates parent checkpoint
- WacGraph.svelte for runtime execution visualization
- Sequential and parallel workflows tested end-to-end

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: WAC v2 bundle cache, globalThis ctx sharing, description optional

- Disable bun bundle caching for WAC v2 scripts (wrapper needs
  windmill-client from node_modules, not available in bundle mode)
- Use Reflect.set/get(globalThis, "__wmill_wf_ctx") to share workflow
  context across dual module instances (wrapper vs user script)
- Never-resolving thenable for non-matching steps in child job mode
  prevents Promise.all race conditions
- Make description field optional in NewScript API (defaults to "")

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add step() primitive for inline checkpointed steps

step() executes a function inline (no child job) and persists the result
to the checkpoint. On replay, the cached value is returned — ensuring
deterministic behavior for non-deterministic operations like Date.now()
or Math.random().

- TypeScript: step(name, fn) — executes inline, throws StepSuspend with
  mode "inline_checkpoint" to persist before continuing
- Rust: InlineCheckpoint variant in WacOutput, saves to checkpoint and
  resets running=false for immediate re-pickup (no zombie wait)
- Shared step counter between task() and step() via _allocKey()

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add Python WAC v2 support with task(), step(), workflow()

- Python SDK: WorkflowCtx with _executing_key child mode, _alloc_key
  shared counter, _run_inline_step for step(), _execute_directly and
  _never_resolve for child mode, step() async function
- Python executor: WAC v2 detection, checkpoint.json writing, WAC
  wrapper.py generation calling _run_workflow(), post-execution hook
  into shared handle_wac_v2_output()
- Make handle_wac_v2_output pub so both bun and python executors share
  the same dispatch/suspend/inline-checkpoint logic
- 17 Python tests covering dispatch, replay, parallel, conditional,
  inline checkpoint, and child mode

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* chore: update sqlx prepared queries

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: WacGraph Tooltip→Popover, simplify wacToFlow parsers

- Fix type error: Tooltip doesn't accept text snippet, use Popover
- Extract shared helpers for task matching and block collection
- Replace linear tasks.find() with Map lookups
- Remove mutable module-level counter

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: Box::pin WAC v2 output handler to prevent stack overflow

handle_python_job's async state machine was too large when combined
with handle_wac_v2_output. Box::pin heap-allocates the future.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: merge WAC v1 and v2 task decorators to preserve backward compat

The v2 @task decorator was shadowing the v1 one, breaking WAC v1
scripts that rely on HTTP-based dispatch via /workflow_as_code/ API.

The merged decorator handles three modes:
- v2: inside @workflow context → checkpoint/replay dispatch
- v1: WM_JOB_ID set, no @workflow → HTTP API dispatch + wait_job
- standalone: no Windmill env → execute function body directly

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: skip no_main_func detection for WAC v2 scripts in TS and Python parsers

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: prevent empty/noop dispatch causing infinite requeue loop

- Validate steps.len() > 0 in WAC dispatch handler (issue 3)
- Replace noop StepSuspend throw with never-resolving promise so it
  can't reach the backend as an empty dispatch (issue 4)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: Python task wrapper now converts positional args to kwargs in v2 mode

Previously only **kwargs were passed to _next_step(), silently dropping
positional arguments. Extract shared _merge_args() helper used by both
v1 and v2 paths.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: replace unwrap() with proper error propagation in WAC arg serialization

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: add workspace_id filter to v2_job queries in WAC dispatch

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: prevent race condition in WAC child dispatch

Restructure dispatch to save checkpoint + suspend parent + seed child
checkpoints in a single transaction BEFORE pushing child jobs. This
ensures a fast child can't complete before the parent is suspended.

Also wrap InlineCheckpoint save + running reset in a transaction to
prevent corrupted state on crash.

Use ULID for pre-generated child job IDs (consistent with rest of API).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: include step key and child job ID in WAC error propagation

Move step_key lookup before the success check so failed child errors
include which task failed, the child job ID, and the original error.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* docs: document WAC determinism contract and step dispatch semantics

- Document that workflow functions must be deterministic across replays
- Document that WacStepDispatch.script/args are metadata, not dispatch targets
- Add comments on counter-based key allocation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: tighten WAC v2 detection to reduce false positives

Replace naive substring matching with line-aware checks that skip
comments and look for specific patterns:
- TS: import from "windmill-client" containing workflow/task
- Python: @workflow and @task decorators with wmill import

Extracted shared helpers in wac_executor.rs used by both executors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: show failed steps in WacGraph when workflow completes with errors

When flowDone is true and a pending step isn't in completedSteps,
mark it as 'failed' instead of 'running'. The failed state CSS and
XCircle icon were already defined but never triggered.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: unsuspend and fail parent when WAC child push fails

Previously if a child push failed mid-batch, the parent remained
suspended with suspend = num_steps but fewer children, hanging until
the 14-day timeout. Now the push loop catches errors and unsuspends
the parent before returning the error.

Also adds source hash validation: if the script content changes between
replays, the job fails with a clear error instead of silently feeding
stale checkpoint data into wrong steps.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: clear suspend_until when unsuspending WAC parent

Set suspend_until = NULL alongside suspend = 0 in both the child
failure and all-children-complete paths, so the parent doesn't rely
on subtle pull query invariants to be re-picked-up.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* test: add exhaustive edge case tests for WAC v2 SDK

fix: make TS task wrapper non-async to fix unawaited task flush

The async wrapper caused microtask-based thenable auto-resolution that
fired .then() and threw StepSuspend before _flushPending() could capture
unawaited steps — making the flush mechanism completely broken. Now the
thenable is returned directly without async wrapping. Backward compatible
with v1 (all code paths still return awaitables).

Tests added (59 TS + 66 Python) covering: full sequential lifecycle,
step after parallel, parallel after parallel, conditional on step result,
empty/single-task workflows, 10+ steps, falsy value preservation, inline
steps, mixed step/task, unawaited flush, child mode with parallel,
key determinism, large parallel groups, and complex mixed patterns.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: atomic checkpoint updates to prevent parallel child race condition

Replace read-modify-write pattern in handle_wac_child_completion with
atomic SQL operations:
- completed_steps merged via jsonb_set(... || jsonb_build_object(...))
  so concurrent children on different workers don't overwrite each other
- suspend counter decremented atomically with RETURNING to determine
  "all done" condition (instead of checking completed_steps in memory)
- suspend_until cleared in the same atomic decrement statement

Before this fix, two parallel children completing simultaneously could
both load the same checkpoint, each add their step, and save — the
second write would overwrite the first, silently losing a child result
and leaving the parent suspended forever.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: cancel already-pushed children on partial WAC dispatch failure

When pushing child jobs sequentially, if pushing child N fails, children
1..N-1 are already running. Previously the error handler only unsuspended
the parent, leaving orphaned children that would complete and corrupt the
checkpoint state (decrementing suspend on an already-unsuspended parent,
potentially causing duplicate step execution on re-run).

Now on partial failure:
1. Cancel all already-pushed children (prevents them from completing
   and corrupting checkpoint state)
2. Clear pending_steps from checkpoint (so parent doesn't think
   children are outstanding on re-run)
3. Then unsuspend parent (so the error propagates)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: skip WAC duration write and child check for non-WAC parents

The duration write to workflow_as_code_status was running for every
non-flow child with a parent (error handlers, success handlers,
run_script children), even though it was only intended for WAC jobs.

Add WHERE workflow_as_code_status IS NOT NULL to skip non-WAC parents
entirely. Piggyback RETURNING pending_steps.job_ids on the same query
so WAC v2 child completion needs zero extra DB round-trips on the
success path.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: seed child checkpoint in same transaction as push

The child checkpoint insert was happening before the child job was
pushed, violating the FK constraint on v2_job_status. Move it into
the push transaction so the job row exists and the child can't be
picked up before its checkpoint is ready.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: set running=false when WAC parent suspends for child dispatch

The parent job kept running=true after suspending, so workers wouldn't
pick it up when children completed and suspend reached 0. The parent
only advanced when the zombie job detector reset it (~90s). Now the
dispatch suspend sets running=false so the parent is immediately
eligible for pickup.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: WAC parent suspend/unsuspend lifecycle

Keep running=true when suspending the parent so the normal pull query
(WHERE running=false) never picks it up. Keep suspend_until non-null
when decrementing suspend to 0 so the suspended pull query
(WHERE suspend_until IS NOT NULL AND suspend<=0) picks it up.

Previously: setting running=false caused infinite restart loops because
the normal pull query has no suspend check and would immediately re-pick
the parent. Clearing suspend_until on the last child prevented the
suspended pull from ever seeing it, requiring the 90s zombie detector.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add approval primitive, flow child completion, timeline fixes for WAC v2

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add error propagation, task options, sleep, and parallel for WAC v2

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* test: fix python SDK tests to use name-based keys and add new test coverage

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: address WAC v2 review findings (sleep timing, error marker, atomicity)

- Fix sleep using suspend=1 instead of 0 to enforce actual delay
- Add approval/sleep resume injection to Python executor
- Fix TS SDK concurrency_limit mapping (was reading wrong property)
- Namespace error marker as __wmill_error to avoid user data collision
- Wrap child completion SQL in transaction for atomicity
- Decrement suspend even when step key is missing (prevents hang)
- Expand TASK_RE to handle export const, let, var, generics
- Validate step key uniqueness before dispatch
- Log warning on checkpoint deserialization failure
- Remove unimplemented delete_after_use from SDKs
- Add TaskError exception class to Python SDK with diagnostic context
- Fix extra positional args handling and add functools.wraps
- Improve getParamNames to handle typed/destructured params

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* sqlx

* sqlx

* test: add WAC v1 e2e integration tests for TS and Python

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: revert fake test versions in typescript-client

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: remove unused WacGraph component and strip wacToFlow to isWorkflowAsCode

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: extract shared approval/sleep resume logic into wac_executor

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 19:39:24 +00:00

833 lines
30 KiB
Rust

use base64::{engine, Engine as _};
use chrono::Datelike;
use core::fmt::Write;
use futures::future::BoxFuture;
use futures::{FutureExt, StreamExt, TryStreamExt};
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use reqwest::{Client, Response};
use serde_json::{json, value::RawValue, Value};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use windmill_common::error::to_anyhow;
use windmill_common::worker::{Connection, SqlResultCollectionStrategy};
use windmill_object_store::convert_json_line_stream;
use windmill_common::{error::Error, worker::to_raw_value};
use windmill_parser_sql::{
parse_db_resource, parse_s3_mode, parse_snowflake_sig, parse_sql_blocks,
};
use windmill_queue::{CanceledBy, MiniPulledJob, HTTP_CLIENT};
use serde::{Deserialize, Serialize};
use crate::common::{build_args_values, get_reserved_variables};
use crate::common::{
build_http_client, resolve_job_timeout, s3_mode_args_to_worker_data, OccupancyMetrics,
S3ModeWorkerData,
};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use windmill_common::client::AuthedClient;
#[derive(Serialize)]
struct Claims {
iss: String,
sub: String,
iat: i64,
exp: i64,
}
#[derive(Deserialize)]
struct SnowflakeDatabase {
account_identifier: String,
public_key: Option<String>,
private_key: Option<String>,
username: Option<String>,
database: Option<String>,
schema: Option<String>,
warehouse: Option<String>,
role: Option<String>,
}
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct SnowflakeResponse {
data: Vec<Vec<Value>>,
resultSetMetaData: SnowflakeResultSetMetaData,
statementHandle: String,
}
#[derive(Deserialize, Debug)]
struct SnowflakeDataOnlyResponse {
data: Vec<Vec<Value>>,
}
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct SnowflakeResultSetMetaData {
numRows: i64,
rowType: Vec<SnowflakeRowType>,
partitionInfo: Vec<Value>,
}
#[derive(Deserialize, Debug)]
struct SnowflakeRowType {
name: String,
r#type: String,
}
#[allow(non_snake_case)]
#[derive(Deserialize)]
struct SnowflakeError {
message: String,
}
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
struct SnowflakeAsyncResponse {
statement_handle: String,
}
async fn poll_snowflake_async_query(
http_client: &Client,
account_identifier: &str,
statement_handle: &str,
token: &str,
token_is_keypair: bool,
deadline: std::time::Instant,
) -> windmill_common::error::Result<SnowflakeResponse> {
let url = format!(
"https://{}.snowflakecomputing.com/api/v2/statements/{}",
account_identifier.to_uppercase(),
statement_handle
);
loop {
if std::time::Instant::now() > deadline {
return Err(Error::ExecutionErr(
"Snowflake query timed out while polling for results".to_string(),
));
}
let mut request = http_client.get(&url).bearer_auth(token);
if token_is_keypair {
request = request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
}
let response = request.send().await.map_err(|e| {
Error::ExecutionErr(format!("Could not poll Snowflake status: {:?}", e))
})?;
let status = response.status();
let body = response
.text()
.await
.map_err(|e| Error::ExecutionErr(format!("error reading poll response body: {}", e)))?;
tracing::debug!(
"Snowflake poll response status: {}, body: {}",
status,
&body[..body.len().min(500)]
);
if status == reqwest::StatusCode::ACCEPTED {
// Still running, wait and poll again
tracing::info!("Snowflake query still running, polling again in 1s...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
continue;
}
if !status.is_success() {
return Err(Error::ExecutionErr(format!(
"Snowflake poll returned error status {}: {}",
status,
&body[..body.len().min(500)]
)));
}
// Query completed, parse the response
let response: SnowflakeResponse = serde_json::from_str(&body).map_err(|e| {
Error::ExecutionErr(format!(
"error decoding poll response: {}. Status: {}. Body preview: {}",
e,
status,
&body[..body.len().min(500)]
))
})?;
return Ok(response);
}
}
async fn handle_snowflake_result(
result: Result<Response, reqwest::Error>,
) -> windmill_common::error::Result<Response> {
match result {
Ok(response) => match response.error_for_status_ref() {
Ok(_) => Ok(response),
Err(e) => {
let resp = response.text().await.unwrap_or("".to_string());
match serde_json::from_str::<SnowflakeError>(&resp) {
Ok(sf_err) => return Err(Error::ExecutionErr(sf_err.message)),
Err(_) => return Err(Error::ExecutionErr(e.to_string())),
}
}
},
Err(e) => Err(Error::ExecutionErr(format!(
"Could not send request: {:?}",
e
))),
}
}
fn do_snowflake_inner<'a>(
query: &'a str,
job_args: &HashMap<String, Value>,
mut body: serde_json::Map<String, Value>,
account_identifier: &'a str,
token: &'a str,
token_is_keypair: bool,
column_order: Option<&'a mut Option<Vec<String>>>,
skip_collect: bool,
first_row_only: bool,
http_client: &'a Client,
s3: Option<S3ModeWorkerData>,
reserved_variables: &HashMap<String, String>,
deadline: std::time::Instant,
) -> windmill_common::error::Result<BoxFuture<'a, windmill_common::error::Result<Vec<Box<RawValue>>>>>
{
let sig = parse_snowflake_sig(&query)
.map_err(|x| Error::ExecutionErr(x.to_string()))?
.args;
let (query, args_to_skip) =
&sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args, reserved_variables)?;
body.insert("statement".to_string(), json!(query));
let mut bindings = serde_json::Map::new();
let mut i = 1;
for arg in &sig {
if args_to_skip.contains(&arg.name) {
continue;
}
let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string());
let arg_v = job_args.get(&arg.name).cloned().unwrap_or(json!(""));
let snowflake_v = convert_typ_val(arg_t, arg_v);
bindings.insert(i.to_string(), snowflake_v);
i += 1;
}
if i > 1 {
body.insert("bindings".to_string(), json!(bindings));
}
let result_f = async move {
let mut request = http_client
.post(format!(
"https://{}.snowflakecomputing.com/api/v2/statements/",
account_identifier.to_uppercase()
))
.bearer_auth(token)
.json(&body);
if token_is_keypair {
request = request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
}
let result = request.send().await;
if skip_collect {
// Still need to handle async (202) responses even when not collecting results
let raw_response = handle_snowflake_result(result).await?;
let status = raw_response.status();
if status == reqwest::StatusCode::ACCEPTED {
let body = raw_response.text().await.map_err(|e| {
Error::ExecutionErr(format!("error reading response body: {}", e))
})?;
let async_resp: SnowflakeAsyncResponse =
serde_json::from_str(&body).map_err(|e| {
Error::ExecutionErr(format!(
"error decoding async response: {}. Body preview: {}",
e,
&body[..body.len().min(500)]
))
})?;
tracing::info!(
"Snowflake statement running asynchronously, polling for completion (handle: {})",
async_resp.statement_handle
);
// Poll until complete, but discard the results
poll_snowflake_async_query(
http_client,
account_identifier,
&async_resp.statement_handle,
token,
token_is_keypair,
deadline,
)
.await?;
}
Ok(vec![])
} else {
// Handle both sync (200) and async (202) responses
let raw_response = handle_snowflake_result(result).await?;
let status = raw_response.status();
let body = raw_response
.text()
.await
.map_err(|e| Error::ExecutionErr(format!("error reading response body: {}", e)))?;
tracing::debug!(
"Snowflake response status: {}, body: {}",
status,
&body[..body.len().min(1000)]
);
let response = if status == reqwest::StatusCode::ACCEPTED {
// Async execution - need to poll for results
let async_resp: SnowflakeAsyncResponse =
serde_json::from_str(&body).map_err(|e| {
Error::ExecutionErr(format!(
"error decoding async response: {}. Body preview: {}",
e,
&body[..body.len().min(500)]
))
})?;
tracing::info!(
"Snowflake query running asynchronously, polling for results (handle: {})",
async_resp.statement_handle
);
poll_snowflake_async_query(
http_client,
account_identifier,
&async_resp.statement_handle,
token,
token_is_keypair,
deadline,
)
.await?
} else {
// Sync execution - parse directly
serde_json::from_str::<SnowflakeResponse>(&body).map_err(|e| {
Error::ExecutionErr(format!(
"error decoding response body: {}. Status: {}. Body preview: {}",
e,
status,
&body[..body.len().min(500)]
))
})?
};
if s3.is_none() && response.resultSetMetaData.numRows > 10000 {
return Err(Error::ExecutionErr(
"More than 10000 rows were requested, use LIMIT 10000 to limit the number of rows or use S3 streaming for larger datasets: https://windmill.dev/docs/core_concepts/sql_to_s3_streaming"
.to_string(),
));
}
if let Some(column_order) = column_order {
*column_order = Some(
response
.resultSetMetaData
.rowType
.iter()
.map(|x| x.name.clone())
.collect::<Vec<String>>(),
);
}
// Clones are because, in s3 mode, reqwest::Body::wrap_stream requires the stream to be
// 'static even though it doesn't make sense to be in our case since the request is
// awaited and the stream is fully read before the function returns.
// Turns out it is a real pain to trick the compiler, even using unsafe
let cloned_account_identifier: String = account_identifier.to_string();
let cloned_token = token.to_string();
let rows_stream = async_stream::stream! {
for row in response.data {
yield Ok::<Vec<Value>, windmill_common::error::Error>(row);
}
if response.resultSetMetaData.partitionInfo.len() > 1 {
for idx in 1..response.resultSetMetaData.partitionInfo.len() {
let url = format!(
"https://{}.snowflakecomputing.com/api/v2/statements/{}",
cloned_account_identifier.to_uppercase(),
response.statementHandle
);
let mut request = HTTP_CLIENT
.get(url)
.bearer_auth(cloned_token.as_str())
.query(&[("partition", idx.to_string())]);
if token_is_keypair {
request =
request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
}
let result = request.send().await;
let raw_response = match handle_snowflake_result(result).await {
Ok(r) => r,
Err(e) => {
yield Err(e);
return;
}
};
let status = raw_response.status();
let body = match raw_response.text().await {
Ok(b) => b,
Err(e) => {
yield Err(Error::ExecutionErr(format!("error reading partition response: {}", e)));
return;
}
};
// Handle async (202) response for partition fetch
let partition_data: SnowflakeDataOnlyResponse = if status == reqwest::StatusCode::ACCEPTED {
// Poll until complete - partition fetches should be fast, but handle async just in case
let mut poll_body = body;
loop {
if std::time::Instant::now() > deadline {
yield Err(Error::ExecutionErr(
"Snowflake partition fetch timed out while polling".to_string(),
));
return;
}
let async_resp: SnowflakeAsyncResponse = match serde_json::from_str(&poll_body) {
Ok(r) => r,
Err(e) => {
yield Err(Error::ExecutionErr(format!(
"error decoding async partition response: {}",
e
)));
return;
}
};
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
let poll_url = format!(
"https://{}.snowflakecomputing.com/api/v2/statements/{}",
cloned_account_identifier.to_uppercase(),
async_resp.statement_handle
);
let mut poll_request = HTTP_CLIENT.get(&poll_url).bearer_auth(cloned_token.as_str());
if token_is_keypair {
poll_request = poll_request.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT");
}
let poll_response = match poll_request.send().await {
Ok(r) => r,
Err(e) => {
yield Err(Error::ExecutionErr(format!("partition poll error: {:?}", e)));
return;
}
};
let poll_status = poll_response.status();
poll_body = match poll_response.text().await {
Ok(b) => b,
Err(e) => {
yield Err(Error::ExecutionErr(format!("error reading partition poll response: {}", e)));
return;
}
};
if poll_status == reqwest::StatusCode::ACCEPTED {
continue;
}
if !poll_status.is_success() {
yield Err(Error::ExecutionErr(format!(
"partition poll returned error: {}",
&poll_body[..poll_body.len().min(500)]
)));
return;
}
match serde_json::from_str(&poll_body) {
Ok(r) => break r,
Err(e) => {
yield Err(Error::ExecutionErr(format!(
"error decoding partition poll response: {}",
e
)));
return;
}
}
}
} else {
match serde_json::from_str(&body) {
Ok(r) => r,
Err(e) => {
yield Err(Error::ExecutionErr(format!(
"error decoding partition response: {}. Body: {}",
e,
&body[..body.len().min(500)]
)));
return;
}
}
};
for row in partition_data.data {
yield Ok(row);
}
}
}
};
let rows_stream = rows_stream.map_ok(move |row| {
let mut row_map = serde_json::Map::new();
row.iter()
.zip(response.resultSetMetaData.rowType.iter())
.for_each(|(val, row_type)| {
row_map.insert(row_type.name.clone(), parse_val(&val, &row_type.r#type));
});
row_map
});
let rows_stream = rows_stream.take(if first_row_only { 1 } else { usize::MAX });
if let Some(s3) = s3 {
let rows_stream =
rows_stream.map(|r| serde_json::value::to_value(&r?).map_err(to_anyhow));
let stream = convert_json_line_stream(rows_stream.boxed(), s3.format).await?;
s3.upload(stream.boxed()).await?;
Ok(vec![to_raw_value(&s3.to_return_s3_obj())])
} else {
let rows = rows_stream
.collect::<Vec<_>>()
.await
.into_iter()
.map(|x| x.map(|v| to_raw_value(&v)))
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
}
}
};
Ok(result_f.boxed())
}
pub async fn do_snowflake(
job: &MiniPulledJob,
client: &AuthedClient,
query: &str,
conn: &Connection,
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
worker_name: &str,
column_order: &mut Option<Vec<String>>,
occupancy_metrics: &mut OccupancyMetrics,
parent_runnable_path: Option<String>,
) -> windmill_common::error::Result<Box<RawValue>> {
let snowflake_args = build_args_values(job, client, conn).await?;
let inline_db_res_path = parse_db_resource(&query);
let s3 = parse_s3_mode(&query)?.map(|s3| s3_mode_args_to_worker_data(s3, client.clone(), job));
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
)
.await?,
)
} else {
snowflake_args.get("database").cloned()
};
let database = if let Some(ref db) = db_arg {
serde_json::from_value::<SnowflakeDatabase>(db.clone())
.map_err(|e| Error::ExecutionErr(e.to_string()))?
} else {
return Err(Error::BadRequest("Missing database argument".to_string()));
};
let annotations = windmill_common::worker::SqlAnnotations::parse(query);
let collection_strategy = if annotations.return_last_result {
SqlResultCollectionStrategy::LastStatementAllRows
} else {
annotations.result_collection
};
// Check if the token is present in db_arg and use it if available
let (token, token_is_keypair) = if let Some(token) = db_arg
.as_ref()
.and_then(|db| db.get("token"))
.and_then(|t| t.as_str())
.filter(|t| !t.is_empty())
{
tracing::debug!("Using oauth token from db_arg");
(token.to_string(), false)
} else {
tracing::debug!("Generating new oauth token");
let qualified_username = format!(
"{}.{}",
database.account_identifier.split('.').next().unwrap_or(""),
database.username.as_deref().unwrap_or("")
)
.to_uppercase();
let public_key = match database.public_key.as_deref() {
Some(key) => pem::parse(key.as_bytes()).map_err(|e| {
Error::ExecutionErr(format!("Failed to parse public key: {}", e.to_string()))
})?,
None => return Err(Error::ExecutionErr("Public key is missing".to_string())),
};
let mut public_key_hash = Sha256::new();
public_key_hash.update(public_key.contents());
let public_key_fp = engine::general_purpose::STANDARD.encode(public_key_hash.finalize());
let iss = format!("{}.SHA256:{}", qualified_username, public_key_fp);
let claims = Claims {
iss: iss,
sub: qualified_username,
iat: chrono::Utc::now().timestamp(),
exp: (chrono::Utc::now() + chrono::Duration::try_hours(1).unwrap()).timestamp(),
};
let private_key = match database.private_key.as_deref() {
Some(key) => EncodingKey::from_rsa_pem(key.as_bytes()).map_err(|e| {
Error::ExecutionErr(format!("Failed to parse private key: {}", e.to_string()))
})?,
None => return Err(Error::ExecutionErr("Private key is missing".to_string())),
};
(
encode(&Header::new(Algorithm::RS256), &claims, &private_key)
.map_err(|e| Error::ExecutionErr(e.to_string()))?,
true,
)
};
tracing::debug!("Snowflake token: {}", token);
let mut body = serde_json::Map::new();
if database.schema.is_some() {
body.insert(
"schema".to_string(),
json!(database.schema.unwrap().to_uppercase()),
);
}
if database.warehouse.is_some() {
body.insert(
"warehouse".to_string(),
json!(database.warehouse.unwrap().to_uppercase()),
);
}
if database.role.is_some() {
body.insert(
"role".to_string(),
json!(database.role.unwrap().to_uppercase()),
);
}
if database.database.is_some() {
body.insert(
"database".to_string(),
json!(database.database.unwrap().to_uppercase()),
);
}
let timeout = resolve_job_timeout(&conn, &job.workspace_id, job.id, job.timeout)
.await
.0
.as_secs();
body.insert("timeout".to_string(), json!(timeout));
let queries = parse_sql_blocks(query);
let (timeout_duration, _, _) =
resolve_job_timeout(&conn, &job.workspace_id, job.id, job.timeout).await;
let http_client = build_http_client(timeout_duration)?;
let deadline = std::time::Instant::now() + timeout_duration;
let reserved_variables =
get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
let result_f = async move {
let mut results = vec![];
for (i, q) in queries.iter().enumerate() {
let result = do_snowflake_inner(
q,
&snowflake_args,
body.clone(),
&database.account_identifier,
&token,
token_is_keypair,
if i == queries.len() - 1
&& s3.is_none()
&& collection_strategy.collect_last_statement_only(queries.len())
&& !collection_strategy.collect_scalar()
{
Some(column_order)
} else {
None
},
collection_strategy.collect_last_statement_only(queries.len())
&& i < queries.len() - 1,
collection_strategy.collect_first_row_only(),
&http_client,
s3.clone(),
&reserved_variables,
deadline,
)?
.await?;
results.push(result);
}
collection_strategy.collect(results)
};
let r = run_future_with_polling_update_job_poller(
job.id,
job.timeout,
conn,
mem_peak,
canceled_by,
result_f,
worker_name,
&job.workspace_id,
&mut Some(occupancy_metrics),
Box::pin(futures::stream::once(async { 0 })),
)
.await?;
*mem_peak = (r.get().len() / 1000) as i32;
Ok(r)
}
fn convert_typ_val(arg_t: String, arg_v: Value) -> Value {
match arg_t.as_str() {
"date" => {
let arr = arg_v
.as_str()
.unwrap_or("")
.split("T")
.collect::<Vec<&str>>();
json!({
"type": "TEXT",
"value": match arr.as_slice() {
[date, _] => {
json!(date)
}
_ => {
arg_v
}
}
})
}
"time" => {
let arr = arg_v
.as_str()
.unwrap_or("")
.split("T")
.collect::<Vec<&str>>();
json!({
"type": "TEXT",
"value": match arr.as_slice() {
[_, time] => {
json!(time)
}
_ => {
arg_v
}
}
})
}
"binary" => {
// convert base64 to hex as expected by snowflake
let bytes = engine::general_purpose::STANDARD
.decode(arg_v.as_str().unwrap_or(""))
.unwrap_or(vec![]);
let mut hex = String::with_capacity(bytes.len() * 2);
for byte in bytes {
write!(hex, "{:02X}", byte).unwrap_or(());
}
json!({
"type": "TEXT",
"value": hex
})
}
_ => {
let mut v = arg_v;
if !v.is_string() {
// if not string, convert to string for api request
v = json!(v.to_string());
}
json!({
"type": "TEXT", // snowflake infer type from schema
"value": v
})
}
}
}
fn parse_val(value: &Value, typ: &str) -> Value {
let str_value = value.as_str().unwrap_or("").to_string();
let val = match typ.to_lowercase().as_str() {
"boolean" => str_value.parse::<bool>().ok().map(|v| json!(v)),
"real" => str_value.parse::<f64>().ok().map(|v| json!(v)),
"timestamp_ltz" | "timestamp_ntz" => str_value
.parse::<f64>()
.ok()
.map(|v| {
chrono::DateTime::from_timestamp(v.round() as i64, 0)
.map(|d| json!(d.format("%Y-%m-%d %H:%M:%S").to_string()))
})
.flatten(),
"time" => str_value
.parse::<f64>()
.ok()
.map(|v| {
chrono::NaiveTime::from_num_seconds_from_midnight_opt(v.round() as u32, 0)
.map(|d| json!(d.format("%H:%M:%S").to_string()))
})
.flatten(),
"date" => str_value
.parse::<i32>()
.ok()
.map(|v| {
chrono::NaiveDate::from_num_days_from_ce_opt(
v + chrono::NaiveDate::from_ymd_opt(1970, 1, 1)
.unwrap()
.num_days_from_ce(),
)
.map(|d| json!(d.format("%Y-%m-%d").to_string()))
})
.flatten(),
"fixed" | "number" => str_value
.parse::<i64>()
.ok()
.map(|v| json!(v))
.or(str_value.parse::<f64>().ok().map(|v| json!(v))),
_ => Some(value.clone()),
};
if let Some(val) = val {
val
} else {
json!(format!(
"ERR: Could not parse {} argument with value {}",
typ, str_value
))
}
}