mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-25 08:00:59 +00:00
fix: Handle CTEs and local tables in SQL asset parser (#8131)
* Handle CTEs and local tables in SQL asset parser * also handle CREATE VIEW * Update package regex version
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::{BTreeMap, HashSet};
|
||||
|
||||
use sqlparser::{
|
||||
ast::{
|
||||
@@ -45,6 +45,10 @@ struct AssetCollector {
|
||||
var_identifiers: BTreeMap<String, (AssetKind, String)>,
|
||||
// e.g USE dl;
|
||||
currently_used_asset: Option<(AssetKind, String)>,
|
||||
// CTE names in scope (stack for nested queries)
|
||||
cte_name_stack: Vec<HashSet<String>>,
|
||||
// Locally created tables (not attached to an asset)
|
||||
local_table_names: HashSet<String>,
|
||||
}
|
||||
|
||||
impl AssetCollector {
|
||||
@@ -54,9 +58,30 @@ impl AssetCollector {
|
||||
current_access_type_stack: Vec::with_capacity(8),
|
||||
var_identifiers: BTreeMap::new(),
|
||||
currently_used_asset: None,
|
||||
cte_name_stack: Vec::new(),
|
||||
local_table_names: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// If the name resolves to an attached asset, record it. Otherwise, register it as a local
|
||||
/// table/view so that subsequent references are not mistakenly attributed to the active asset.
|
||||
fn track_table_definition(&mut self, name: &ObjectName) {
|
||||
if let Some(asset) = self.get_associated_asset_from_obj_name(name, Some(W)) {
|
||||
self.assets.push(asset);
|
||||
} else if let Some(simple_name) = get_trivial_obj_name(name) {
|
||||
self.local_table_names.insert(simple_name.to_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
fn is_locally_defined(&self, name: &str) -> bool {
|
||||
let name_lower = name.to_lowercase();
|
||||
self.local_table_names.contains(&name_lower)
|
||||
|| self
|
||||
.cte_name_stack
|
||||
.iter()
|
||||
.any(|set| set.contains(&name_lower))
|
||||
}
|
||||
|
||||
// Detect when we do 'a.b' and 'a' is associated with an asset in var_identifiers
|
||||
// Or when we access 'b' and we did USE a;
|
||||
fn get_associated_asset_from_obj_name(
|
||||
@@ -72,6 +97,14 @@ impl AssetCollector {
|
||||
return None;
|
||||
}
|
||||
|
||||
if name.0.len() == 1 {
|
||||
if let Some(ident) = name.0.first().and_then(|id| id.as_ident()) {
|
||||
if self.is_locally_defined(&ident.value) {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if name.0.len() == 1 || name.0.len() == 2 {
|
||||
if name
|
||||
.0
|
||||
@@ -452,6 +485,7 @@ impl Visitor for AssetCollector {
|
||||
) -> std::ops::ControlFlow<Self::Break> {
|
||||
match statement {
|
||||
sqlparser::ast::Statement::Query(q) => {
|
||||
self.cte_name_stack.push(collect_cte_names(q));
|
||||
if let Some(select) = q.body.as_select() {
|
||||
// First, handle table references (adds table-level assets)
|
||||
for t in &select.from {
|
||||
@@ -612,17 +646,11 @@ impl Visitor for AssetCollector {
|
||||
}
|
||||
|
||||
sqlparser::ast::Statement::CreateTable(create_table) => {
|
||||
if let Some(asset) =
|
||||
self.get_associated_asset_from_obj_name(&create_table.name, Some(W))
|
||||
{
|
||||
self.assets.push(asset);
|
||||
}
|
||||
self.track_table_definition(&create_table.name);
|
||||
}
|
||||
|
||||
sqlparser::ast::Statement::CreateView { name, .. } => {
|
||||
if let Some(asset) = self.get_associated_asset_from_obj_name(name, Some(W)) {
|
||||
self.assets.push(asset);
|
||||
}
|
||||
self.track_table_definition(name);
|
||||
}
|
||||
|
||||
sqlparser::ast::Statement::Copy { target: CopyTarget::File { filename }, .. } => {
|
||||
@@ -672,16 +700,20 @@ impl Visitor for AssetCollector {
|
||||
|
||||
fn post_visit_statement(
|
||||
&mut self,
|
||||
_statement: &sqlparser::ast::Statement,
|
||||
statement: &sqlparser::ast::Statement,
|
||||
) -> std::ops::ControlFlow<Self::Break> {
|
||||
if matches!(statement, sqlparser::ast::Statement::Query(_)) {
|
||||
self.cte_name_stack.pop();
|
||||
}
|
||||
std::ops::ControlFlow::Continue(())
|
||||
}
|
||||
|
||||
fn pre_visit_query(
|
||||
&mut self,
|
||||
_query: &sqlparser::ast::Query,
|
||||
query: &sqlparser::ast::Query,
|
||||
) -> std::ops::ControlFlow<Self::Break> {
|
||||
self.current_access_type_stack.push(R);
|
||||
self.cte_name_stack.push(collect_cte_names(query));
|
||||
std::ops::ControlFlow::Continue(())
|
||||
}
|
||||
|
||||
@@ -690,12 +722,22 @@ impl Visitor for AssetCollector {
|
||||
_query: &sqlparser::ast::Query,
|
||||
) -> std::ops::ControlFlow<Self::Break> {
|
||||
self.current_access_type_stack.pop();
|
||||
self.cte_name_stack.pop();
|
||||
std::ops::ControlFlow::Continue(())
|
||||
}
|
||||
|
||||
// We do not use pre_visit_relation because we cannot know if an ObjectName is a table or a function
|
||||
}
|
||||
|
||||
fn collect_cte_names(query: &sqlparser::ast::Query) -> HashSet<String> {
|
||||
query.with.as_ref().map_or_else(HashSet::new, |with| {
|
||||
with.cte_tables
|
||||
.iter()
|
||||
.map(|cte| cte.alias.name.value.to_lowercase())
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
fn is_read_fn(fname: &str) -> bool {
|
||||
fname.eq_ignore_ascii_case("read_parquet")
|
||||
|| fname.eq_ignore_ascii_case("read_csv")
|
||||
@@ -1509,6 +1551,235 @@ mod tests {
|
||||
assert!(result[0].columns.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_cte_not_treated_as_asset() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
WITH tmp AS (SELECT 1 AS x)
|
||||
SELECT * FROM tmp;
|
||||
SELECT * FROM real_table;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/real_table".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_cte_scope_does_not_leak() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
WITH tmp AS (SELECT 1) SELECT * FROM tmp;
|
||||
SELECT * FROM tmp;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/tmp".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_multiple_ctes() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
WITH cte1 AS (SELECT 1), cte2 AS (SELECT 2)
|
||||
SELECT * FROM cte1 JOIN cte2 ON true;
|
||||
SELECT * FROM real_table;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/real_table".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_local_create_table_overrides_asset() {
|
||||
let input = r#"
|
||||
CREATE TABLE local_tbl (id INT);
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
SELECT * FROM local_tbl;
|
||||
SELECT * FROM asset_table;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/asset_table".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_create_table_with_use_is_still_asset() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake' AS dl; USE dl;
|
||||
CREATE TABLE friends (
|
||||
name text,
|
||||
age int
|
||||
);
|
||||
INSERT INTO friends VALUES ($name, $age);
|
||||
SELECT * FROM friends;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "main/friends".to_string(),
|
||||
access_type: Some(RW),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_local_create_view_overrides_asset() {
|
||||
let input = r#"
|
||||
CREATE VIEW my_view AS SELECT 1;
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
SELECT * FROM my_view;
|
||||
SELECT * FROM asset_table;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/asset_table".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_create_view_with_use_is_still_asset() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
CREATE VIEW my_view AS SELECT 1;
|
||||
SELECT * FROM my_view;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/my_view".to_string(),
|
||||
access_type: Some(RW),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_cte_mixed_with_asset_tables() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
WITH tmp AS (SELECT 1 AS x)
|
||||
SELECT * FROM tmp JOIN real_table ON true;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/real_table".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_local_table_insert_and_select() {
|
||||
let input = r#"
|
||||
CREATE TABLE staging (id INT, val TEXT);
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
INSERT INTO staging VALUES (1, 'a');
|
||||
SELECT * FROM staging;
|
||||
INSERT INTO real_table VALUES (2, 'b');
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/real_table".to_string(),
|
||||
access_type: Some(W),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_qualified_ref_bypasses_local() {
|
||||
// Even if 'tbl' is local, 'dl.tbl' is an explicit asset reference
|
||||
let input = r#"
|
||||
CREATE TABLE tbl (id INT);
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
SELECT * FROM dl.tbl;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl/tbl".to_string(),
|
||||
access_type: Some(R),
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_cte_case_insensitive() {
|
||||
let input = r#"
|
||||
ATTACH 'ducklake://my_dl' AS dl;
|
||||
USE dl;
|
||||
WITH MyTable AS (SELECT 1)
|
||||
SELECT * FROM mytable;
|
||||
"#;
|
||||
let s = parse_assets(input).map(|s| s.assets);
|
||||
assert_eq!(
|
||||
s.map_err(|e| e.to_string()),
|
||||
Ok(vec![ParseAssetsResult {
|
||||
kind: AssetKind::Ducklake,
|
||||
path: "my_dl".to_string(),
|
||||
access_type: None,
|
||||
columns: None
|
||||
},])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sql_asset_parser_s3_read_csv_columns() {
|
||||
let input = r#"
|
||||
|
||||
Generated
+4
-4
@@ -80,7 +80,7 @@
|
||||
"windmill-parser-wasm-nu": "1.510.1",
|
||||
"windmill-parser-wasm-php": "1.574.1",
|
||||
"windmill-parser-wasm-py": "^1.628.3",
|
||||
"windmill-parser-wasm-regex": "1.639.0",
|
||||
"windmill-parser-wasm-regex": "1.646.0",
|
||||
"windmill-parser-wasm-ruby": "1.526.1",
|
||||
"windmill-parser-wasm-rust": "1.558.1",
|
||||
"windmill-parser-wasm-ts": "1.623.1",
|
||||
@@ -14684,9 +14684,9 @@
|
||||
"integrity": "sha512-TlluqknZpg8cZ+A3m6JFLPseY2PpKtDsxdj26fAnCUzKPtse8TxQR+n0dwC80rfW5TwdWSulvNGRDgcNuf7CTw=="
|
||||
},
|
||||
"node_modules/windmill-parser-wasm-regex": {
|
||||
"version": "1.639.0",
|
||||
"resolved": "https://registry.npmjs.org/windmill-parser-wasm-regex/-/windmill-parser-wasm-regex-1.639.0.tgz",
|
||||
"integrity": "sha512-qvYM4sYxB6M0xrqwBljS2fWqOMk6rp++60TRltJnzZDzVaWQrKjTGwNMmfepGAIWy1OGVKp0SCVERhe2P+O6tQ=="
|
||||
"version": "1.646.0",
|
||||
"resolved": "https://registry.npmjs.org/windmill-parser-wasm-regex/-/windmill-parser-wasm-regex-1.646.0.tgz",
|
||||
"integrity": "sha512-sRDQBX9ML3VIcW9r1Ug8SOKkHog6uKBwCTE6OADU3K8o8C45NiJufMGCeKtBWoF/Ki08IEj+0h4dd6QEIKpeIQ=="
|
||||
},
|
||||
"node_modules/windmill-parser-wasm-ruby": {
|
||||
"version": "1.526.1",
|
||||
|
||||
@@ -152,7 +152,7 @@
|
||||
"windmill-parser-wasm-nu": "1.510.1",
|
||||
"windmill-parser-wasm-php": "1.574.1",
|
||||
"windmill-parser-wasm-py": "^1.628.3",
|
||||
"windmill-parser-wasm-regex": "1.639.0",
|
||||
"windmill-parser-wasm-regex": "1.646.0",
|
||||
"windmill-parser-wasm-ruby": "1.526.1",
|
||||
"windmill-parser-wasm-rust": "1.558.1",
|
||||
"windmill-parser-wasm-ts": "1.623.1",
|
||||
|
||||
Reference in New Issue
Block a user