mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-18 16:02:10 +00:00
feat: Ducklake typechecker (#8118)
* Typedchecked ducklake queries * Display script preview error as SQL error * Fix duplication * fix replacer * Revert "fix replacer" This reverts commitc5492033c8. * 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 commitae64a8ad27. * 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:
+1
@@ -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"] }
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user