feat: Ducklake typechecker (#8118)

* Typedchecked ducklake queries

* Display script preview error as SQL error

* Fix duplication

* fix replacer

* Revert "fix replacer"

This reverts commit c5492033c8.

* 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 ae64a8ad27.

* 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>
This commit is contained in:
Diego Imbert
2026-03-04 11:46:08 +01:00
committed by GitHub
parent 424ca59dfe
commit 53caecf1da
5 changed files with 423 additions and 89 deletions
+1
View File
@@ -2164,6 +2164,7 @@ version = "0.1.0"
dependencies = [
"chrono",
"duckdb",
"regex",
"rust_decimal",
"serde",
"serde_json",
@@ -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"] }
+219 -36
View File
@@ -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<Vec<PrepareQueryColumnInfo>>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
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<Regex> = 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::<Vec<_>>()
};
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<String, String> {
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<PrepareQueryResult> = 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::<std::result::Result<Vec<_>, _>>()
}) {
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<Vec<String>>), 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<Box<RawValue>>> = vec![];
let mut column_order = None;
@@ -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<Item = &'a str>,
token: &str,
base_internal_url: &str,
w_id: &str,
) -> Result<Box<RawValue>> {
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::<Result<Vec<_>>>()?;
let query_block_list = query_block_list
.iter()
.map(|s| s.as_ptr())
.collect::<Vec<_>>();
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,
+121 -53
View File
@@ -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
}