From 53caecf1da8d76e246178dfb9b86d330f0ec52fd Mon Sep 17 00:00:00 2001 From: Diego Imbert <70353967+diegoimbert@users.noreply.github.com> Date: Wed, 4 Mar 2026 11:46:08 +0100 Subject: [PATCH] feat: Ducklake typechecker (#8118) * Typedchecked ducklake queries * Display script preview error as SQL error * Fix duplication * fix replacer * Revert "fix replacer" This reverts commit c5492033c850cabc8bf18a50151c089b83cd6826. * Don't recompile regex every call * nit OOB * avoid potential panic * Apply suggestions from code review Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> * safety throw * Update backend/windmill-worker/src/duckdb_executor.rs Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> * Try catch individual chunks in prepareDatatableQueries Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> * format * nit comment * Revert "Try catch individual chunks in prepareDatatableQueries" This reverts commit ae64a8ad27deb7e5ddda10163c6db04a428827f2. * Correct try catch * better error messages * nit unused variable * comment * handle non describable queries * npm i --------- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> --- .../windmill-duckdb-ffi-internal/Cargo.lock | 1 + .../windmill-duckdb-ffi-internal/Cargo.toml | 1 + .../windmill-duckdb-ffi-internal/src/lib.rs | 255 +++++++++++++++--- .../windmill-worker/src/duckdb_executor.rs | 81 ++++++ frontend/src/lib/infer.svelte.ts | 174 ++++++++---- 5 files changed, 423 insertions(+), 89 deletions(-) diff --git a/backend/windmill-duckdb-ffi-internal/Cargo.lock b/backend/windmill-duckdb-ffi-internal/Cargo.lock index 07e428c633..559196a3c2 100644 --- a/backend/windmill-duckdb-ffi-internal/Cargo.lock +++ b/backend/windmill-duckdb-ffi-internal/Cargo.lock @@ -2164,6 +2164,7 @@ version = "0.1.0" dependencies = [ "chrono", "duckdb", + "regex", "rust_decimal", "serde", "serde_json", diff --git a/backend/windmill-duckdb-ffi-internal/Cargo.toml b/backend/windmill-duckdb-ffi-internal/Cargo.toml index 7043b33ee5..7eb6869ab9 100644 --- a/backend/windmill-duckdb-ffi-internal/Cargo.toml +++ b/backend/windmill-duckdb-ffi-internal/Cargo.toml @@ -6,6 +6,7 @@ edition = "2024" [dependencies] chrono = "0.4.41" duckdb = { version = "1.4.4", features = ["bundled"] } +regex = "1" rust_decimal = "1.37.2" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "^1", features = ["preserve_order", "raw_value"] } diff --git a/backend/windmill-duckdb-ffi-internal/src/lib.rs b/backend/windmill-duckdb-ffi-internal/src/lib.rs index c5c819e60b..2701319d6e 100644 --- a/backend/windmill-duckdb-ffi-internal/src/lib.rs +++ b/backend/windmill-duckdb-ffi-internal/src/lib.rs @@ -1,12 +1,14 @@ use std::{ collections::HashMap, - ffi::{CStr, CString, c_char, c_uint}, + ffi::{c_char, c_uint, CStr, CString}, ptr::null_mut, + sync::LazyLock, }; -use duckdb::{Row, core::LogicalTypeId, params_from_iter, types::TimeUnit}; -use rust_decimal::{Decimal, prelude::FromPrimitive}; -use serde::Deserialize; +use duckdb::{core::LogicalTypeId, params_from_iter, types::TimeUnit, Row}; +use regex::Regex; +use rust_decimal::{prelude::FromPrimitive, Decimal}; +use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; #[derive(Deserialize, Clone, Debug, PartialEq, Default)] @@ -96,6 +98,218 @@ pub extern "C" fn run_duckdb_ffi( }) } +#[derive(Serialize, Debug)] +struct PrepareQueryColumnInfo { + name: String, + #[serde(rename = "type")] + type_name: String, +} + +#[derive(Serialize, Debug)] +struct PrepareQueryResult { + #[serde(skip_serializing_if = "Option::is_none")] + columns: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +fn is_setup_statement(query: &str) -> bool { + let trimmed = query.trim_start(); + let upper = trimmed.to_uppercase(); + upper.starts_with("ATTACH") + || upper.starts_with("USE") + || upper.starts_with("INSTALL") + || upper.starts_with("LOAD") + || upper.starts_with("SET") + || upper.starts_with("RESET") + || upper.starts_with("CREATE OR REPLACE SECRET") + || upper.starts_with("CREATE SECRET") +} + +/// Returns true if the query is expected to return a result set and can be wrapped with DESCRIBE. +fn is_describable_query(query: &str) -> bool { + let trimmed = query.trim_start(); + let upper = trimmed.to_uppercase(); + upper.starts_with("SELECT") + || upper.starts_with("WITH") + || upper.starts_with("VALUES") + || upper.starts_with("TABLE") + || upper.starts_with("FROM") +} + +static PARAM_RE: LazyLock = LazyLock::new(|| Regex::new(r"\$\d+").expect("invalid regex")); + +fn replace_params_with_null(query: &str) -> String { + PARAM_RE.replace_all(query, "NULL").to_string() +} + +#[unsafe(no_mangle)] +pub extern "C" fn prepare_duckdb_ffi( + query_block_list: *const *const c_char, + query_block_list_count: usize, + token: *const c_char, + base_internal_url: *const c_char, + w_id: *const c_char, +) -> *mut c_char { + let r = match convert_prepare_args( + query_block_list, + query_block_list_count, + token, + base_internal_url, + w_id, + ) + .and_then(|(query_block_list, token, base_internal_url, w_id)| { + prepare_duckdb_internal(query_block_list, token, base_internal_url, w_id) + }) { + Ok(result) => result, + Err(err) => { + let err = serde_json::to_string(&err) + .unwrap_or_else(|_| "Unknown error in duckdb ffi lib".to_string()); + format!("ERROR {}", err) + } + }; + + CString::new(r).map(|s| s.into_raw()).unwrap_or_else(|e| { + println!("Failed to allocate error string in duckdb ffi lib: {:?}", e); + null_mut() + }) +} + +fn setup_duckdb_connection( + conn: &duckdb::Connection, + token: &str, + base_internal_url: &str, + w_id: &str, +) -> Result<(), String> { + let (s3_access_key, s3_secret_key) = token.rsplit_once('.').unwrap_or(("", token)); + let (s3_endpoint_ssl, s3_endpoint) = base_internal_url + .split_once("://") + .unwrap_or(("http", &base_internal_url)); + let s3_endpoint_ssl = s3_endpoint_ssl == "https"; + + conn.execute_batch(&format!( + "INSTALL httpfs; LOAD httpfs; + INSTALL azure; LOAD azure; + CREATE OR REPLACE SECRET s3_secret ( + TYPE s3, + PROVIDER config, + KEY_ID '{s3_access_key}', + SECRET '{s3_secret_key}', + ENDPOINT '{s3_endpoint}/api/w/{w_id}/s3_proxy', + URL_STYLE path, + USE_SSL {s3_endpoint_ssl} + ); + CREATE OR REPLACE SECRET gcs_secret ( + TYPE gcs, + KEY_ID '{s3_access_key}', + SECRET '{s3_secret_key}', + ENDPOINT '{s3_endpoint}/api/w/{w_id}/s3_proxy', + USE_SSL {s3_endpoint_ssl} + ); + ", + )) + .map_err(|e| format!("Error setting up S3 secret: {}", e.to_string())) +} + +fn convert_prepare_args<'a>( + query_block_list: *const *const c_char, + query_block_list_count: usize, + token: *const c_char, + base_internal_url: *const c_char, + w_id: *const c_char, +) -> Result<(Vec<&'a str>, &'a str, &'a str, &'a str), String> { + let query_block_list = unsafe { + std::slice::from_raw_parts(query_block_list, query_block_list_count) + .iter() + .map(|q| { + CStr::from_ptr(*q).to_str().unwrap_or_else(|e| { + println!( + "Invalid query_block string pointer in duckdb ffi: {}", + e.to_string() + ); + "Invalid query_block string pointer in duckdb ffi" + }) + }) + .collect::>() + }; + let token = unsafe { CStr::from_ptr(token) } + .to_str() + .map_err(|e| format!("Invalid token string: {}", e.to_string()))?; + let base_internal_url = unsafe { CStr::from_ptr(base_internal_url) } + .to_str() + .map_err(|e| format!("Invalid base_internal_url string: {}", e.to_string()))?; + let w_id = unsafe { CStr::from_ptr(w_id) } + .to_str() + .map_err(|e| format!("Invalid w_id string: {}", e.to_string()))?; + Ok((query_block_list, token, base_internal_url, w_id)) +} + +fn prepare_duckdb_internal( + query_block_list: Vec<&str>, + token: &str, + base_internal_url: &str, + w_id: &str, +) -> Result { + let conn = duckdb::Connection::open_in_memory().map_err(|e| e.to_string())?; + + setup_duckdb_connection(&conn, token, base_internal_url, w_id)?; + + let mut results: Vec = vec![]; + + // IMPORTANT: Setup statements (ATTACH, USE, INSTALL, etc.) are executed but intentionally + // do not produce a PrepareQueryResult entry. The frontend prepends these as connection setup + // before the actual user queries, and mapPrepareResults expects results.length to equal the + // number of user queries (not setup statements). If a new setup-like statement is added to + // the connection flow (e.g. in setup_duckdb_connection or transform_attach_ducklake) without + // also being caught by is_setup_statement, the result count will mismatch and the frontend + // will throw. + for query_block in &query_block_list { + if is_setup_statement(query_block) { + conn.execute_batch(query_block) + .map_err(|e| format!("Error executing setup statement: {}", e.to_string()))?; + continue; + } + + let modified_query = replace_params_with_null(query_block); + // Validate the query parses correctly by preparing it + if let Err(e) = conn.prepare(&modified_query) { + results.push(PrepareQueryResult { columns: None, error: Some(e.to_string()) }); + continue; + } + + // DESCRIBE only works on queries that return result sets (SELECT, WITH, VALUES, TABLE, + // FROM). For non-returning statements (INSERT, UPDATE, DELETE, CREATE, DROP, ALTER, etc.) + // we skip DESCRIBE and assume no columns. + if !is_describable_query(&modified_query) { + results.push(PrepareQueryResult { columns: Some(vec![]), error: None }); + continue; + } + + // Note: We have to use a DESCRIBE statement and cannot simply use the + // methods returned by .prepare() because they panic if the statement was + // not executed at least once (which we specifically do not want to do). + let describe_query = format!("DESCRIBE {}", modified_query); + match conn.prepare(&describe_query).and_then(|mut stmt| { + let rows = stmt.query_map([], |row| { + Ok(PrepareQueryColumnInfo { + name: row.get::<_, String>(0)?, + type_name: row.get::<_, String>(1)?, + }) + })?; + rows.collect::, _>>() + }) { + Ok(columns) => { + results.push(PrepareQueryResult { columns: Some(columns), error: None }); + } + Err(e) => { + results.push(PrepareQueryResult { columns: None, error: Some(e.to_string()) }); + } + } + } + + serde_json::to_string(&results).map_err(|e| e.to_string()) +} + fn convert_args<'a>( query_block_list: *const *const c_char, query_block_list_count: usize, @@ -170,38 +384,7 @@ fn run_duckdb_internal<'a>( ) -> Result<(String, Option>), String> { let conn = duckdb::Connection::open_in_memory().map_err(|e| e.to_string())?; - let (s3_access_key, s3_secret_key) = token.split_at(token.rfind('.').unwrap_or(0)); - let s3_secret_key = &s3_secret_key[1..]; - let (s3_endpoint_ssl, s3_endpoint) = base_internal_url - .split_once("://") - .unwrap_or(("http", &base_internal_url)); - let s3_endpoint_ssl = match s3_endpoint_ssl { - "https" => true, - _ => false, - }; - - conn.execute_batch(&format!( - "INSTALL httpfs; LOAD httpfs; - INSTALL azure; LOAD azure; - CREATE OR REPLACE SECRET s3_secret ( - TYPE s3, - PROVIDER config, - KEY_ID '{s3_access_key}', - SECRET '{s3_secret_key}', - ENDPOINT '{s3_endpoint}/api/w/{w_id}/s3_proxy', - URL_STYLE path, - USE_SSL {s3_endpoint_ssl} - ); - CREATE OR REPLACE SECRET gcs_secret ( - TYPE gcs, - KEY_ID '{s3_access_key}', - SECRET '{s3_secret_key}', - ENDPOINT '{s3_endpoint}/api/w/{w_id}/s3_proxy', - USE_SSL {s3_endpoint_ssl} - ); - ", - )) - .map_err(|e| format!("Error setting up S3 secret: {}", e.to_string()))?; + setup_duckdb_connection(&conn, token, base_internal_url, w_id)?; let mut results: Vec>> = vec![]; let mut column_order = None; diff --git a/backend/windmill-worker/src/duckdb_executor.rs b/backend/windmill-worker/src/duckdb_executor.rs index 73136e1cc0..45e2f647a5 100644 --- a/backend/windmill-worker/src/duckdb_executor.rs +++ b/backend/windmill-worker/src/duckdb_executor.rs @@ -161,6 +161,22 @@ pub async fn do_duckdb( let base_internal_url = client.base_internal_url.clone(); let w_id = job.workspace_id.clone(); + if annotations.prepare { + let result = tokio::task::spawn_blocking(move || { + prepare_duckdb_ffi_safe( + query_block_list.iter().map(String::as_str), + &token, + &base_internal_url, + &w_id, + ) + }) + .await + .map_err(|e| Error::from(to_anyhow(e))) + .and_then(|r| r)?; + + return Ok(result); + } + let result = tokio::task::spawn_blocking(move || { run_duckdb_ffi_safe( query_block_list.iter().map(String::as_str), @@ -248,6 +264,18 @@ struct DuckDbFfiLib { collect_first_row_only: bool, ) -> *mut c_char, >, + prepare_duckdb_ffi: Option< + Symbol< + 'static, + unsafe extern "C" fn( + query_block_list: *const *const c_char, + query_block_list_count: usize, + token: *const c_char, + base_internal_url: *const c_char, + w_id: *const c_char, + ) -> *mut c_char, + >, + >, free_cstr: Symbol<'static, unsafe extern "C" fn(string: *mut c_char) -> ()>, } @@ -307,8 +335,11 @@ impl DuckDbFfiLib { } } + let prepare_duckdb_ffi = unsafe { lib.get(b"prepare_duckdb_ffi").ok() }; + Ok(DuckDbFfiLib { run_duckdb_ffi: unsafe { lib.get(b"run_duckdb_ffi").map_err(to_anyhow)? }, + prepare_duckdb_ffi, free_cstr: unsafe { lib.get(b"free_cstr").map_err(to_anyhow)? }, }) } @@ -388,6 +419,56 @@ fn run_duckdb_ffi_safe<'a>( } } +fn prepare_duckdb_ffi_safe<'a>( + query_block_list: impl Iterator, + token: &str, + base_internal_url: &str, + w_id: &str, +) -> Result> { + let query_block_list = query_block_list + .map(|s| { + CString::new(s).map_err(|e| { + Error::ExecutionErr(format!("Failed CString conversion: {}", e.to_string())) + }) + }) + .collect::>>()?; + let query_block_list = query_block_list + .iter() + .map(|s| s.as_ptr()) + .collect::>(); + + let token = CString::new(token).map_err(to_anyhow)?; + let base_internal_url = CString::new(base_internal_url).map_err(to_anyhow)?; + let w_id = CString::new(w_id).map_err(to_anyhow)?; + + let lib = DuckDbFfiLib::get_singleton()?; + let prepare_fn = lib.prepare_duckdb_ffi.as_ref().ok_or_else(|| { + Error::InternalErr( + "prepare_duckdb_ffi not available in duckdb ffi library. Please update to the latest windmill_duckdb_ffi_lib.".to_string(), + ) + })?; + let free_cstr = &lib.free_cstr; + + let result_str = unsafe { + let ptr = prepare_fn( + query_block_list.as_ptr(), + query_block_list.len(), + token.as_ptr(), + base_internal_url.as_ptr(), + w_id.as_ptr(), + ); + let str = CStr::from_ptr(ptr).to_string_lossy().to_string(); + free_cstr(ptr); + str + }; + + if result_str.starts_with("ERROR") { + Err(Error::ExecutionErr(result_str[6..].to_string())) + } else { + Ok(serde_json::value::RawValue::from_string(result_str).map_err(to_anyhow)?) + } +} + struct ParsedAttachDbResource<'a> { resource_path: &'a str, name: &'a str, diff --git a/frontend/src/lib/infer.svelte.ts b/frontend/src/lib/infer.svelte.ts index 419020631c..cbdfa95505 100644 --- a/frontend/src/lib/infer.svelte.ts +++ b/frontend/src/lib/infer.svelte.ts @@ -4,6 +4,13 @@ import { ChangeOnDeepInequality, MapResource } from './svelte5Utils.svelte' import { sqlDataTypeToJsTypeHeuristic } from './components/apps/components/display/dbtable/utils' import { chunkBy, clone, getQueryStmtCountHeuristic } from './utils' +function extractErrorMessage(e: unknown): string { + if (e != null && typeof e === 'object' && 'body' in e) { + return (e as any).body?.error?.message ?? JSON.stringify(e) + } + return e instanceof Error ? e.message : JSON.stringify(e) +} + function computeQueryKey(query: InferAssetsSqlQueryDetails, workspace?: string) { return `${query.source_kind}::${query.source_name}::${query.source_schema}::${workspace}::${query.query_string}` } @@ -21,66 +28,26 @@ export function usePreparedAssetSqlQueries( ), async (toFetch) => { let queries = Object.entries(clone(toFetch)) - // We only support preparing datatable source kinds for now. - queries = queries.filter(([_, q]) => q.source_kind === 'datatable') + queries = queries.filter( + ([_, q]) => q.source_kind === 'datatable' || q.source_kind === 'ducklake' + ) // We only support preparing single-statement queries for now. queries = queries.filter(([_, q]) => getQueryStmtCountHeuristic(q.query_string) === 1) if (!queries?.length) return {} - try { - // We chunk by source_name to minimize the number of requests. - // For example if we have 10 queries on the same data table, - // we can prepare them all with a single script. - queries.sort((a, b) => a[1].source_name.localeCompare(b[1].source_name)) - let results = ( - await Promise.all( - chunkBy(queries, ([key, q]) => q.source_name).map(async (chunk) => { - console.log( - 'Preparing chunk of queries:', - chunk.map(([_, q]) => q) - ) - let queryContent = chunk - .flatMap(([key, q]) => [ - q.source_schema ? `SET search_path TO ${q.source_schema};` : 'RESET search_path;', - q.query_string + (q.query_string.trim().endsWith(';') ? '' : ';') - ]) - .join('\n') - queryContent = - '-- prepare\n--result_collection=all_statements_first_row\n' + queryContent + let datatableQueries = queries.filter(([_, q]) => q.source_kind === 'datatable') + let ducklakeQueries = queries.filter(([_, q]) => q.source_kind === 'ducklake') - let res = (await JobService.runScriptPreviewAndWaitResult({ - workspace: getWorkspace()!, - requestBody: { - language: 'postgresql', - content: queryContent, - args: { database: `datatable://${chunk[0][1]?.source_name}` } - } - })) as { error?: string; columns?: { name: string; type: string }[] }[] + let allResults: [string, PreparedAssetsSqlQuery][] = [] - console.log('Prepared query content:', res) - - let res2: [string, PreparedAssetsSqlQuery][] = res.map((r, i) => [ - chunk[i][0], - r.columns - ? { - columns: Object.fromEntries( - r.columns.map(({ name, type }) => [ - name, - sqlDataTypeToJsTypeHeuristic(type) - ]) - ) - } - : { error: r.error ?? "Couldn't prepare query " } - ]) - return res2 - }) - ) - ).flat() - - return Object.fromEntries(results) - } catch (e) { - throw e + if (datatableQueries.length) { + allResults.push(...(await prepareDatatableQueries(datatableQueries, getWorkspace))) } + if (ducklakeQueries.length) { + allResults.push(...(await prepareDucklakeQueries(ducklakeQueries, getWorkspace))) + } + + return Object.fromEntries(allResults) } ) @@ -96,3 +63,104 @@ export function usePreparedAssetSqlQueries( } } } + +type QueryEntry = [string, InferAssetsSqlQueryDetails] + +function mapPrepareResults( + res: { error?: string; columns?: { name: string; type: string }[] }[], + chunk: QueryEntry[] +): [string, PreparedAssetsSqlQuery][] { + if (res.length !== chunk.length) { + throw new Error(`Prepare results count mismatch: got ${res.length}, expected ${chunk.length}`) + } + return res.map((r, i) => [ + chunk[i]?.[0], + r.columns + ? { + columns: Object.fromEntries( + r.columns.map(({ name, type: t }) => [name, sqlDataTypeToJsTypeHeuristic(t)]) + ) + } + : { error: r.error ?? "Couldn't prepare query " } + ]) +} + +async function prepareDatatableQueries( + queries: QueryEntry[], + getWorkspace: () => string | undefined +): Promise<[string, PreparedAssetsSqlQuery][]> { + queries.sort((a, b) => a[1].source_name.localeCompare(b[1].source_name)) + let results = ( + await Promise.all( + chunkBy(queries, ([_, q]) => q.source_name).map(async (chunk) => { + let queryContent = chunk + .flatMap(([_, q]) => [ + q.source_schema ? `SET search_path TO ${q.source_schema};` : 'RESET search_path;', + q.query_string + (q.query_string.trim().endsWith(';') ? '' : ';') + ]) + .join('\n') + queryContent = '-- prepare\n--result_collection=all_statements_first_row\n' + queryContent + + try { + let res = (await JobService.runScriptPreviewAndWaitResult({ + workspace: getWorkspace()!, + requestBody: { + language: 'postgresql', + content: queryContent, + args: { database: `datatable://${chunk[0][1]?.source_name}` } + } + })) as { error?: string; columns?: { name: string; type: string }[] }[] + + return mapPrepareResults(res, chunk) + } catch (e) { + const error = extractErrorMessage(e) + return chunk.map(([key]) => [key, { error }] as [string, PreparedAssetsSqlQuery]) + } + }) + ) + ).flat() + return results +} + +async function prepareDucklakeQueries( + queries: QueryEntry[], + getWorkspace: () => string | undefined +): Promise<[string, PreparedAssetsSqlQuery][]> { + queries.sort((a, b) => a[1].source_name.localeCompare(b[1].source_name)) + let results = ( + await Promise.all( + chunkBy(queries, ([_, q]) => `${q.source_name}::${q.source_schema ?? ''}`).map( + async (chunk) => { + let sourceName = chunk[0][1].source_name + let sourceSchema = chunk[0][1].source_schema + let attachSetup = `ATTACH 'ducklake://${sourceName}' AS dl;\n` + attachSetup += sourceSchema ? `USE dl.${sourceSchema};\n` : `USE dl;\n` + + let queryContent = chunk + .map(([_, q]) => q.query_string + (q.query_string.trim().endsWith(';') ? '' : ';')) + .join('\n') + queryContent = + '-- prepare\n--result_collection=all_statements_first_row\n' + + attachSetup + + queryContent + + try { + let res = (await JobService.runScriptPreviewAndWaitResult({ + workspace: getWorkspace()!, + requestBody: { + language: 'duckdb', + content: queryContent, + args: {} + } + })) as { error?: string; columns?: { name: string; type: string }[] }[] + return mapPrepareResults(res, chunk) + } catch (e) { + const error = extractErrorMessage(e) + return chunk.map(([key]) => [key, { error }] as [string, PreparedAssetsSqlQuery]) + } + } + ) + ) + ).flat() + return results +}