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:
Diego Imbert
2026-02-28 08:04:19 +01:00
committed by GitHub
parent 257734b9ab
commit 095505136c
3 changed files with 287 additions and 16 deletions
@@ -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#"
+4 -4
View File
@@ -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",
+1 -1
View File
@@ -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",