use polyglot_sql::{
analyze_query, scope::SourceKind, AnalyzeQueryOptions, DialectType, ProjectionNullability,
QueryShape, ReferenceConfidence, SetOperationBranchRole, TransformKind, ValidationSchema,
};
use serde_json::json;
#[test]
fn analyze_query_preserves_cast_types_through_cte_chains() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{"name": "raw_orders", "columns": [{"name": "amount", "type": "VARCHAR"}]}]
}))
.unwrap();
for (dialect, expected) in [
(DialectType::Snowflake, "INT"),
(DialectType::DuckDB, "INT"),
(DialectType::PostgreSQL, "INT"),
(DialectType::BigQuery, "INT64"),
] {
for sql in [
"WITH transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders), final AS (SELECT amount FROM transformed) SELECT amount FROM final",
"WITH transformed AS (SELECT CAST(amount AS INTEGER) AS n FROM raw_orders), final AS (SELECT n FROM transformed) SELECT n FROM final",
"WITH transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders), final AS (SELECT transformed.amount FROM transformed) SELECT final.amount FROM final",
"WITH transformed(n) AS (SELECT CAST(amount AS INTEGER) FROM raw_orders), final(m) AS (SELECT n FROM transformed) SELECT m FROM final",
"WITH transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders), final AS (SELECT * FROM transformed) SELECT * FROM final",
"WITH transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders) SELECT amount FROM (SELECT amount FROM transformed) AS final",
"WITH unused AS (SELECT CAST(amount AS DATE) AS amount FROM raw_orders), transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders) SELECT amount FROM transformed",
] {
for with_schema in [false, true] {
let analysis = analyze_query(sql, AnalyzeQueryOptions {
dialect,
schema: with_schema.then(|| schema.clone()),
}).unwrap();
let fact = &analysis.projections[0];
assert_eq!(fact.type_hint.as_deref(), Some(expected), "{dialect:?}, schema={with_schema}: {sql}");
assert_eq!(fact.transform_kind, TransformKind::Direct, "{sql}");
assert_eq!(fact.cast_type, None, "{sql}");
assert!(fact.upstream.iter().any(|reference| {
reference.table.as_deref().is_some_and(|table| table.eq_ignore_ascii_case("raw_orders"))
&& reference.column.eq_ignore_ascii_case("amount")
}), "{sql}: {:?}", fact.upstream);
}
}
}
}
#[test]
fn analyze_query_keeps_scalar_subquery_types_in_their_own_scope() {
let schema: ValidationSchema = serde_json::from_value(json!({"tables": [
{"name":"t1", "columns":[{"name":"id","type":"INT"}]},
{"name":"t2", "columns":[{"name":"id","type":"INT"},{"name":"val","type":"DOUBLE"}]}
]}))
.unwrap();
for (sql, expected) in [
("SELECT (SELECT val FROM t2 WHERE t2.id = t1.id) AS v FROM t1", "DOUBLE"),
("SELECT (SELECT id) AS v FROM t1", "INT"),
("SELECT (SELECT (SELECT id)) AS v FROM t1", "INT"),
("SELECT (WITH c AS (SELECT id AS n) SELECT n FROM c) AS v FROM t1", "INT"),
("SELECT (WITH c AS (SELECT CAST(val AS VARCHAR) AS n FROM t2) SELECT n FROM c) AS v FROM t1", "TEXT"),
("SELECT b.val FROM t1 JOIN LATERAL (SELECT val FROM t2 WHERE t2.id = t1.id) AS b ON TRUE", "DOUBLE"),
("SELECT b.id FROM t1, LATERAL (SELECT id) AS b", "INT"),
("WITH unused AS (SELECT 1 AS n) SELECT (SELECT n) AS v FROM t1", "UNKNOWN"),
] {
let analysis = analyze_query(sql, AnalyzeQueryOptions {
dialect: DialectType::DuckDB, schema: Some(schema.clone()),
}).unwrap();
assert_eq!(analysis.projections[0].type_hint.as_deref(), Some(expected), "{sql}");
}
}
#[test]
fn analyze_query_infers_the_selected_expression_not_an_arbitrary_upstream_cast() {
for (sql, expected) in [
("WITH a AS (SELECT CAST('1' AS INT) AS n), b AS (SELECT CAST(n AS VARCHAR) AS n FROM a) SELECT n FROM b", "TEXT"),
("WITH a AS (SELECT CAST('1' AS INT) AS n), b AS (SELECT n + 0.5 AS n FROM a) SELECT n FROM b", "DOUBLE"),
("WITH a AS (SELECT CAST('1' AS INT) AS n), b AS (WITH a AS (SELECT CAST('2024-01-01' AS DATE) AS n) SELECT n FROM a) SELECT n FROM b", "DATE"),
("WITH a AS (SELECT CAST('1' AS INT) AS n), b AS (SELECT n FROM a UNION ALL SELECT n FROM a) SELECT n FROM b", "INT"),
] {
let analysis = analyze_query(sql, AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
..Default::default()
}).unwrap();
assert_eq!(analysis.projections[0].type_hint.as_deref(), Some(expected), "{sql}");
}
}
#[test]
fn schema_aware_lineage_preserves_cte_cast_nodes() {
use polyglot_sql::traversal::ExpressionWalk;
let sql = "WITH transformed AS (SELECT CAST(amount AS INTEGER) AS amount FROM raw_orders), final AS (SELECT amount FROM transformed) SELECT amount FROM final";
let expression = polyglot_sql::parse_one(sql, DialectType::Snowflake).unwrap();
let schema = polyglot_sql::mapping_schema_from_validation_schema_with_dialect(
&serde_json::from_value(json!({"tables": [{"name": "raw_orders", "columns": [{"name": "amount", "type": "VARCHAR"}]}]})).unwrap(),
DialectType::Snowflake,
);
let node = polyglot_sql::lineage::lineage_with_schema(
"amount",
&expression,
Some(&schema),
Some(DialectType::Snowflake),
false,
)
.unwrap();
let nodes: Vec<_> = node.walk().collect();
assert_eq!(nodes.len(), 4);
assert!(nodes.iter().any(|node| node
.expression
.contains(|e| matches!(e, polyglot_sql::Expression::Cast(_)))));
assert_eq!(nodes[1].source_kind, SourceKind::Cte);
assert_eq!(nodes[2].source_kind, SourceKind::Cte);
assert_eq!(nodes[3].source_kind, SourceKind::Table);
}
fn column_use_analysis(sql: &str, with_schema: bool) -> polyglot_sql::QueryAnalysis {
analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: with_schema.then(schema),
},
)
.unwrap()
}
#[test]
fn analyze_query_reports_filter_uses_without_changing_projection_lineage() {
use polyglot_sql::ColumnUseContext;
for with_schema in [false, true] {
let sql = "SELECT o.id FROM orders AS o WHERE o.amount > 10";
let analysis = column_use_analysis(sql, with_schema);
assert_eq!(analysis.projections[0].upstream.len(), 1);
assert_eq!(analysis.projections[0].upstream[0].column, "id");
assert_eq!(analysis.column_uses.len(), 1);
let fact = &analysis.column_uses[0];
assert_eq!(fact.context, ColumnUseContext::Filter);
assert_eq!(fact.scope_path, "root");
assert_eq!(fact.expression_path, "where_clause.this");
assert_eq!(fact.expression_sql, "o.amount > 10");
assert!(fact.span.is_none());
let reference = &fact.references[0];
assert_eq!(reference.reference.source_name.as_deref(), Some("orders"));
assert_eq!(reference.reference.source_alias.as_deref(), Some("o"));
assert_eq!(reference.reference.column, "amount");
assert_eq!(
reference.reference.confidence,
ReferenceConfidence::Resolved
);
let span = reference.span.unwrap();
assert_eq!(
sql.chars()
.skip(span.start)
.take(span.end - span.start)
.collect::<String>(),
"o.amount"
);
}
assert!(column_use_analysis("SELECT 1", false)
.column_uses
.is_empty());
}
#[test]
fn analyze_query_column_uses_group_conditions_and_keep_occurrence_spans() {
use polyglot_sql::ColumnUseContext;
let sql = "SELECT '😀', o.id FROM orders o JOIN customers c ON o.customer_id = c.id AND o.amount > 0 WHERE o.amount > 1 OR o.amount < 0";
let analysis = column_use_analysis(sql, true);
let join = analysis
.column_uses
.iter()
.find(|fact| fact.context == ColumnUseContext::Join)
.unwrap();
assert_eq!(join.expression_path, "joins[0].on");
assert_eq!(join.references.len(), 3);
let filter = analysis
.column_uses
.iter()
.find(|fact| fact.context == ColumnUseContext::Filter)
.unwrap();
assert_eq!(filter.references.len(), 2);
assert_ne!(filter.references[0].span, filter.references[1].span);
for reference in &filter.references {
let span = reference.span.unwrap();
assert_eq!(
sql.chars()
.skip(span.start)
.take(span.end - span.start)
.collect::<String>(),
"o.amount"
);
}
}
#[test]
fn analyze_query_column_uses_cover_group_having_order_aliases_and_ordinals() {
use polyglot_sql::ColumnUseContext;
let analysis = column_use_analysis("SELECT o.customer_id AS customer, SUM(o.amount) AS amount_sum FROM orders o GROUP BY 1 HAVING amount_sum > 10 ORDER BY amount_sum, 1", true);
for (context, expected) in [
(ColumnUseContext::Group, "customer_id"),
(ColumnUseContext::Having, "amount"),
] {
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == context)
.unwrap();
assert_eq!(fact.references[0].reference.column, expected, "{fact:?}");
}
let order: Vec<_> = analysis
.column_uses
.iter()
.filter(|fact| fact.context == ColumnUseContext::Order)
.collect();
assert_eq!(order.len(), 2);
assert_eq!(order[0].references[0].reference.column, "amount");
assert_eq!(order[1].references[0].reference.column, "customer_id");
assert!(order[1].references[0].span.is_none());
}
#[test]
fn analyze_query_column_uses_cover_windows_qualify_and_aggregate_filters() {
use polyglot_sql::ColumnUseContext;
let analysis = column_use_analysis("SELECT o.id, ROW_NUMBER() OVER w AS rn, SUM(o.amount) FILTER (WHERE o.amount > 0) OVER (PARTITION BY o.customer_id ORDER BY o.id) AS total FROM orders o WINDOW w AS (PARTITION BY o.customer_id ORDER BY o.id) QUALIFY rn = 1 ORDER BY o.id", true);
for context in [
ColumnUseContext::WindowPartition,
ColumnUseContext::WindowOrder,
ColumnUseContext::Qualify,
ColumnUseContext::Filter,
ColumnUseContext::Order,
] {
let facts: Vec<_> = analysis
.column_uses
.iter()
.filter(|fact| fact.context == context)
.collect();
assert!(
!facts.is_empty(),
"missing {context:?}: {:?}",
analysis.column_uses
);
assert!(
facts.iter().all(|fact| !fact.references.is_empty()),
"{facts:?}"
);
}
}
#[test]
fn analyze_query_column_uses_resolve_chained_ctes_and_correlated_scopes() {
use polyglot_sql::ColumnUseContext;
let sql = "WITH base AS (SELECT id, amount FROM orders), paid AS (SELECT id, amount * 2 AS amount FROM base WHERE amount > 0) SELECT p.id FROM paid p WHERE p.amount > 10 AND EXISTS (SELECT 1 FROM customers c WHERE c.id = p.id)";
for with_schema in [false, true] {
let analysis = column_use_analysis(sql, with_schema);
let root = analysis
.column_uses
.iter()
.find(|fact| fact.scope_path == "root" && fact.context == ColumnUseContext::Filter)
.unwrap();
assert_eq!(
root.references.len(),
1,
"subquery references belong to their own scope"
);
assert_eq!(
root.references[0].reference.table.as_deref(),
Some("orders"),
"{root:?}"
);
assert_eq!(root.references[0].reference.column, "amount");
let correlated = analysis
.column_uses
.iter()
.find(|fact| fact.scope_path == "root.subqueries[0]")
.unwrap();
assert_eq!(correlated.references.len(), 2);
assert_eq!(
correlated.references[1].reference.table.as_deref(),
Some("orders"),
"{correlated:?}"
);
assert!(analysis
.column_uses
.iter()
.any(|fact| fact.scope_path == "root.ctes[1]"));
}
}
#[test]
fn analyze_query_column_uses_do_not_leak_ctes_or_non_lateral_sources() {
let analysis = column_use_analysis("SELECT o.id FROM orders o WHERE EXISTS (WITH local AS (SELECT id FROM customers) SELECT 1 FROM local WHERE local.id = o.id) AND EXISTS (SELECT 1 FROM local WHERE local.id = 1)", true);
let leaked = analysis
.column_uses
.iter()
.find(|fact| fact.scope_path == "root.subqueries[1]")
.unwrap();
assert_ne!(
leaked.references[0].reference.table.as_deref(),
Some("customers")
);
let analysis = column_use_analysis("SELECT o.id FROM orders o JOIN (SELECT id FROM customers WHERE id = o.id) c ON o.id = c.id", false);
let derived = analysis
.column_uses
.iter()
.find(|fact| fact.scope_path == "root.derived[0]")
.unwrap();
assert_eq!(
derived.references[1].reference.confidence,
ReferenceConfidence::Unknown
);
}
#[test]
fn analyze_query_column_uses_conservatively_resolve_partial_and_ambiguous_schemas() {
for (sql, with_schema, expected) in [
(
"SELECT 1 FROM orders o JOIN customers c ON TRUE WHERE id > 0",
true,
ReferenceConfidence::Ambiguous,
),
(
"SELECT 1 FROM orders o JOIN missing m ON TRUE WHERE id > 0",
true,
ReferenceConfidence::Unknown,
),
(
"SELECT 1 FROM orders o JOIN customers c ON TRUE WHERE id > 0",
false,
ReferenceConfidence::Unknown,
),
(
"SELECT 1 FROM orders o WHERE bogus.id > 0",
true,
ReferenceConfidence::Unknown,
),
(
"SELECT 1 FROM orders o WHERE o.missing > 0",
true,
ReferenceConfidence::Unknown,
),
] {
let analysis = column_use_analysis(sql, with_schema);
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::Filter)
.unwrap();
assert_eq!(
fact.references[0].reference.confidence, expected,
"{sql}: {fact:?}"
);
}
}
#[test]
fn analyze_query_column_uses_cover_using_and_natural_joins() {
for sql in [
"SELECT 1 FROM orders JOIN customers USING (id)",
"SELECT 1 FROM orders NATURAL JOIN customers",
] {
let analysis = column_use_analysis(sql, true);
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::Join)
.unwrap();
assert_eq!(fact.references.len(), 2, "{fact:?}");
assert!(fact
.references
.iter()
.all(|reference| reference.reference.column == "id"
&& reference.reference.confidence == ReferenceConfidence::Resolved));
}
let unknown = column_use_analysis("SELECT 1 FROM orders NATURAL JOIN customers", false);
assert_eq!(
unknown.column_uses[0].references[0].reference.confidence,
ReferenceConfidence::Unknown
);
}
#[test]
fn analyze_query_column_uses_identify_nested_filter_branches_and_set_order() {
use polyglot_sql::ColumnUseContext;
let sql = "SELECT id FROM orders EXCEPT (SELECT id FROM customers WHERE id > 0 UNION ALL SELECT id FROM users) ORDER BY id";
let analysis = column_use_analysis(sql, true);
let filter_branches: Vec<_> = analysis
.column_uses
.iter()
.filter(|fact| fact.context == ColumnUseContext::SetOperationFilter)
.collect();
assert_eq!(filter_branches.len(), 2, "{:?}", analysis.column_uses);
assert!(filter_branches
.iter()
.all(|fact| fact.scope_path.starts_with("root.branches[1].branches[")));
let order = analysis
.column_uses
.iter()
.find(|fact| fact.context == ColumnUseContext::Order)
.unwrap();
assert!(!order.references.is_empty());
assert!(
analysis.projections[0]
.upstream
.iter()
.any(|reference| reference.table.as_deref() == Some("customers")),
"existing lineage must remain unchanged"
);
}
#[test]
fn analyze_query_column_uses_are_json_additive_and_deterministic() {
let sql =
"SELECT o.id FROM orders o JOIN customers c USING(id) WHERE o.amount > 0 ORDER BY o.id";
let value = serde_json::to_value(column_use_analysis(sql, true)).unwrap();
assert_eq!(
value,
serde_json::to_value(column_use_analysis(sql, true)).unwrap()
);
assert!(value["columnUses"][0]["references"][0]
.get("reference")
.is_none());
let mut legacy = value;
legacy.as_object_mut().unwrap().remove("columnUses");
let decoded: polyglot_sql::QueryAnalysis = serde_json::from_value(legacy).unwrap();
assert!(decoded.column_uses.is_empty());
}
#[test]
fn analyze_query_column_uses_include_scalar_predicate_inputs_but_not_exists_outputs() {
for predicate in [
"o.id = (SELECT MAX(c.id) FROM customers c)",
"o.id IN (SELECT c.id FROM customers c)",
] {
let analysis =
column_use_analysis(&format!("SELECT 1 FROM orders o WHERE {predicate}"), true);
let root = &analysis.column_uses[0];
assert_eq!(root.references.len(), 2, "{root:?}");
assert!(root
.references
.iter()
.any(|usage| usage.reference.table.as_deref() == Some("customers")));
}
let analysis = column_use_analysis(
"SELECT 1 FROM orders o WHERE EXISTS (SELECT c.name FROM customers c WHERE c.id = o.id)",
true,
);
assert!(analysis.column_uses[0].references.is_empty());
assert_eq!(analysis.column_uses[1].references.len(), 2);
assert!(analysis
.column_uses
.iter()
.flat_map(|fact| &fact.references)
.all(|usage| usage.reference.column != "name"));
}
#[test]
fn analyze_query_column_uses_handle_cte_alias_columns_and_unknown_transitive_references() {
let analysis = column_use_analysis(
"WITH renamed(key) AS (SELECT id FROM orders) SELECT 1 FROM renamed WHERE key > 0",
true,
);
assert_eq!(analysis.column_uses[0].references[0].reference.column, "id");
assert_eq!(
analysis.column_uses[0].references[0]
.reference
.table
.as_deref(),
Some("orders")
);
for sql in [
"WITH invalid AS (SELECT bogus.id AS key FROM orders) SELECT 1 FROM invalid WHERE key > 0",
"WITH invalid AS (SELECT o.missing AS key FROM orders o) SELECT 1 FROM invalid WHERE key > 0",
] {
let analysis = column_use_analysis(sql, true);
assert_eq!(analysis.column_uses[0].references[0].reference.confidence, ReferenceConfidence::Unknown, "{analysis:?}");
}
}
#[test]
fn analyze_query_column_uses_expand_star_filter_inputs_with_schema() {
let analysis = column_use_analysis("SELECT * FROM customers EXCEPT SELECT * FROM users", true);
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::SetOperationFilter)
.unwrap();
assert_eq!(fact.references.len(), 2, "{fact:?}");
assert_eq!(fact.references[0].reference.column, "id");
assert_eq!(fact.references[1].reference.column, "name");
}
#[test]
fn analyze_query_column_uses_resolve_output_aliases_after_star_expansion() {
let analysis = column_use_analysis(
"SELECT o.*, o.amount * 2 AS doubled FROM orders o ORDER BY doubled",
true,
);
let reference = &analysis.column_uses[0].references[0].reference;
assert_eq!(reference.column, "amount");
assert_eq!(reference.table.as_deref(), Some("orders"));
}
#[test]
fn analyze_query_column_uses_reuse_struct_field_resolution() {
let schema: ValidationSchema = serde_json::from_value(json!({"tables":[{"name":"events","columns":[{"name":"payload","type":"STRUCT(active BOOLEAN)"}]}]})).unwrap();
let analysis = analyze_query(
"SELECT 1 FROM events WHERE payload.active = TRUE",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema),
},
)
.unwrap();
let reference = &analysis.column_uses[0].references[0].reference;
assert_eq!(reference.column, "payload");
assert_eq!(reference.table.as_deref(), Some("events"));
assert_eq!(reference.confidence, ReferenceConfidence::Resolved);
}
#[test]
fn analyze_query_column_uses_preserve_transitive_partial_schema_uncertainty() {
let analysis = column_use_analysis(
"WITH c AS (SELECT id FROM orders o JOIN missing m ON TRUE) SELECT 1 FROM c WHERE c.id > 0",
true,
);
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::Filter)
.unwrap();
assert!(
fact.references
.iter()
.all(|usage| usage.reference.confidence == ReferenceConfidence::Unknown),
"{fact:?}"
);
}
#[test]
fn analyze_query_column_uses_cover_supported_qualify_dialects_and_filter_operators() {
for dialect in [
DialectType::DuckDB,
DialectType::BigQuery,
DialectType::Snowflake,
] {
let analysis = analyze_query("SELECT o.id, ROW_NUMBER() OVER (PARTITION BY o.customer_id ORDER BY o.id) AS rn FROM orders o QUALIFY rn = 1", AnalyzeQueryOptions { dialect, schema: Some(schema()) }).unwrap();
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::Qualify)
.unwrap();
assert!(!fact.references.is_empty(), "{dialect:?}: {fact:?}");
assert!(
fact.references
.iter()
.all(|usage| usage.reference.confidence == ReferenceConfidence::Resolved),
"{dialect:?}: {fact:?}"
);
}
for operator in ["EXCEPT", "INTERSECT"] {
let analysis = column_use_analysis(
&format!("SELECT id FROM orders {operator} SELECT id FROM customers"),
true,
);
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == polyglot_sql::ColumnUseContext::SetOperationFilter)
.unwrap();
assert_eq!(fact.scope_path, "root.branches[1]");
assert_eq!(
fact.references[0].reference.table.as_deref(),
Some("customers")
);
}
}
#[test]
fn analyze_query_column_uses_cover_aggregate_order_and_window_frame_inputs() {
let analysis = column_use_analysis("SELECT ARRAY_AGG(o.amount ORDER BY o.id), SUM(o.amount) OVER (ORDER BY o.id ROWS BETWEEN o.customer_id PRECEDING AND CURRENT ROW) FROM orders o", false);
for (context, column) in [
(polyglot_sql::ColumnUseContext::AggregateOrder, "id"),
(polyglot_sql::ColumnUseContext::WindowFrame, "customer_id"),
] {
let fact = analysis
.column_uses
.iter()
.find(|fact| fact.context == context)
.unwrap_or_else(|| panic!("missing {context:?}: {:?}", analysis.column_uses));
assert_eq!(fact.references[0].reference.column, column);
}
}
#[test]
fn analyze_query_column_uses_do_not_resolve_forward_join_aliases() {
let analysis = column_use_analysis(
"SELECT o.id FROM orders o JOIN customers c ON c.id = u.id JOIN users u ON u.id = o.id",
false,
);
let joins: Vec<_> = analysis
.column_uses
.iter()
.filter(|fact| fact.context == polyglot_sql::ColumnUseContext::Join)
.collect();
assert_eq!(
joins[0].references[1].reference.confidence,
ReferenceConfidence::Unknown
);
assert_eq!(
joins[1].references[0].reference.confidence,
ReferenceConfidence::Resolved
);
}
fn schema() -> ValidationSchema {
serde_json::from_value(json!({
"tables": [
{
"name": "users",
"columns": [
{"name": "id", "type": "INT", "primaryKey": true},
{"name": "name", "type": "TEXT", "nullable": false}
]
},
{
"name": "orders",
"columns": [
{"name": "id", "type": "INT", "nullable": false},
{"name": "user_id", "type": "INT"},
{"name": "customer_id", "type": "INT"},
{"name": "amount", "type": "DECIMAL(10,2)", "nullable": true},
{"name": "total", "type": "FLOAT"}
]
},
{
"name": "customers",
"columns": [
{"name": "id", "type": "INT", "nullable": false},
{"name": "name", "type": "TEXT", "nullable": false}
]
},
{
"name": "x",
"columns": [{"name": "a", "type": "INT"}]
},
{
"name": "y",
"columns": [{"name": "b", "type": "INT"}]
}
]
}))
.unwrap()
}
#[test]
fn analyze_query_reports_projection_relations_and_types() {
let analysis = analyze_query(
"SELECT u.id, CAST(o.total AS TEXT) AS total_text, 1 AS one \
FROM users AS u JOIN orders AS o ON u.id = o.user_id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.shape, QueryShape::Select);
assert_eq!(analysis.projections.len(), 3);
assert_eq!(analysis.projections[0].name.as_deref(), Some("id"));
assert_eq!(
analysis.projections[0].transform_kind,
TransformKind::Direct
);
assert_eq!(analysis.projections[1].transform_kind, TransformKind::Cast);
assert_eq!(analysis.projections[1].cast_type.as_deref(), Some("TEXT"));
assert_eq!(
analysis.projections[2].transform_kind,
TransformKind::Constant
);
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "users" && relation.alias.as_deref() == Some("u")));
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "orders" && relation.columns.contains(&"total".into())));
assert!(analysis
.base_tables
.iter()
.any(|relation| relation.name == "orders"));
assert!(analysis
.base_tables
.iter()
.any(|relation| relation.name == "users"));
let total = &analysis.projections[1].upstream;
assert!(total.iter().any(|reference| {
reference.table.as_deref() == Some("orders")
&& reference.column == "total"
&& reference.confidence == ReferenceConfidence::Resolved
}));
}
#[test]
fn analyze_query_duckdb_extract_date_part_types() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "events",
"columns": [{"name": "created_at", "type": "TIMESTAMP"}]
}]
}))
.unwrap();
for (date_part, expected_type) in [
("'year'", "BIGINT"),
("'month'", "BIGINT"),
("'day'", "BIGINT"),
("YEAR", "BIGINT"),
("'second'", "BIGINT"),
("'milliseconds'", "BIGINT"),
("'microseconds'", "BIGINT"),
("'yearweek'", "BIGINT"),
("'epoch'", "DOUBLE"),
("EPOCH", "DOUBLE"),
("'EpOcH'", "DOUBLE"),
("'julian'", "DOUBLE"),
("JULIAN", "DOUBLE"),
("'JuLiAn'", "DOUBLE"),
] {
let sql = format!("SELECT EXTRACT({date_part} FROM created_at) AS extracted FROM events");
let analysis = analyze_query(
&sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema.clone()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert_eq!(analysis.projections[0].name.as_deref(), Some("extracted"));
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some(expected_type),
"unexpected output type for {sql:?}"
);
assert!(analysis.projections[0].upstream.iter().any(|reference| {
reference.table.as_deref() == Some("events") && reference.column == "created_at"
}));
}
}
#[test]
fn analyze_query_duckdb_trim_types_and_lineage() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "records",
"columns": [
{"name": "name", "type": "VARCHAR"},
{"name": "chars", "type": "VARCHAR"},
{"name": "values_json", "type": "JSON"}
]
}]
}))
.unwrap();
for (function, columns) in [
("TRIM(name)", vec!["name"]),
(
r#"TRIM(BOTH '"' FROM values_json->>0)"#,
vec!["values_json"],
),
(
"TRIM(LOWER(TRIM(name)), UPPER(chars))",
vec!["name", "chars"],
),
] {
let sql = format!("SELECT {function} AS normalized FROM records");
for schema in [None, Some(schema.clone())] {
let has_schema = schema.is_some();
let analysis = analyze_query(
&sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema,
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 1, "{sql}");
let projection = &analysis.projections[0];
assert_eq!(projection.name.as_deref(), Some("normalized"), "{sql}");
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"), "{sql}");
if has_schema {
assert_eq!(projection.upstream.len(), columns.len(), "{sql}");
for column in &columns {
assert!(
projection.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("records")
&& reference.column == *column
&& reference.confidence == ReferenceConfidence::Resolved
}),
"{sql}: {column}"
);
}
}
}
}
}
#[test]
fn analyze_query_duckdb_regexp_extract_all_types() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "documents",
"columns": [{"name": "body", "type": "VARCHAR"}]
}]
}))
.unwrap();
for (function, type_hint) in [
("REGEXP_EXTRACT_ALL(body, '[0-9]+')", "TEXT[]"),
(
"REGEXP_EXTRACT_ALL(body, '([a-z])([0-9]+)', 2, 'i')",
"TEXT[]",
),
(
"REGEXP_EXTRACT_ALL(body, '([a-z])([0-9]+)', ['letter', 'number'])",
"STRUCT(letter TEXT, number TEXT)[]",
),
("UNNEST(REGEXP_EXTRACT_ALL(body, '[0-9]+'))", "TEXT"),
(
"UNNEST(REGEXP_EXTRACT_ALL(body, '([a-z])([0-9]+)', ['letter', 'number'], 'i'))",
"STRUCT(letter TEXT, number TEXT)",
),
] {
let sql = format!("SELECT {function} AS matches FROM documents");
for schema in [None, Some(schema.clone())] {
let has_schema = schema.is_some();
let analysis = analyze_query(
&sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema,
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 1, "{sql}");
let projection = &analysis.projections[0];
assert_eq!(projection.name.as_deref(), Some("matches"), "{sql}");
assert_eq!(projection.type_hint.as_deref(), Some(type_hint), "{sql}");
if has_schema {
assert!(
projection.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("documents")
&& reference.column == "body"
&& reference.confidence == ReferenceConfidence::Resolved
}),
"{sql}"
);
}
}
}
}
#[test]
fn analyze_query_duckdb_date_name_types() {
for input_type in ["DATE", "TIMESTAMP", "TIMESTAMPTZ"] {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "events",
"columns": [{"name": "created_at", "type": input_type}]
}]
}))
.unwrap();
for function in ["MONTHNAME", "DAYNAME"] {
let sql = format!("SELECT {function}(created_at) AS label FROM events");
for schema in [None, Some(schema.clone())] {
let has_schema = schema.is_some();
let analysis = analyze_query(
&sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema,
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 1, "{sql}");
let projection = &analysis.projections[0];
assert_eq!(projection.name.as_deref(), Some("label"), "{sql}");
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"), "{sql}");
assert!(
projection.upstream.iter().any(|reference| {
reference.column == "created_at"
&& (!has_schema
|| (reference.table.as_deref() == Some("events")
&& reference.confidence == ReferenceConfidence::Resolved))
}),
"{sql}"
);
}
}
}
}
#[test]
fn analyze_query_duckdb_array_to_string_types() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "events",
"columns": [
{"name": "labels", "type": "VARCHAR[]"},
{"name": "label", "type": "VARCHAR"}
]
}]
}))
.unwrap();
for (function, column) in [
("ARRAY_TO_STRING(labels, ', ')", "labels"),
("ARRAY_TO_STRING(ARRAY_AGG(label), ', ')", "label"),
("ARRAY_TO_STRING_COMMA_DEFAULT(labels)", "labels"),
("ARRAY_TO_STRING_COMMA_DEFAULT(ARRAY_AGG(label))", "label"),
] {
let sql = format!("SELECT {function} AS label_text FROM events");
for schema in [None, Some(schema.clone())] {
let has_schema = schema.is_some();
let analysis = analyze_query(
&sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema,
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert_eq!(analysis.projections.len(), 1, "{sql}");
let projection = &analysis.projections[0];
assert_eq!(projection.name.as_deref(), Some("label_text"), "{sql}");
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"), "{sql}");
if has_schema {
assert!(
projection.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("events")
&& reference.column == column
&& reference.confidence == ReferenceConfidence::Resolved
}),
"{sql}"
);
}
}
}
}
#[test]
fn analyze_query_reports_function_projection_arguments() {
let analysis = analyze_query(
"SELECT DATE_TRUNC('month', created_at) AS bucket FROM events",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(
serde_json::from_value(json!({
"tables": [
{
"name": "events",
"columns": [
{"name": "created_at", "type": "TIMESTAMP"}
]
}
]
}))
.unwrap(),
),
},
)
.unwrap();
let transform_function = analysis.projections[0]
.transform_function
.as_ref()
.expect("expected transform function fact");
assert_eq!(transform_function.name, "DATE_TRUNC");
assert_eq!(transform_function.literal_args, vec!["month"]);
assert_eq!(transform_function.column_args.len(), 1);
assert_eq!(
transform_function.column_args[0].table.as_deref(),
Some("events")
);
assert_eq!(transform_function.column_args[0].column, "created_at");
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some("TIMESTAMP")
);
}
#[test]
fn analyze_query_follows_cte_lineage() {
let analysis = analyze_query(
"WITH base AS (SELECT id FROM users) SELECT id FROM base",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.ctes, vec!["base"]);
assert!(analysis.projections[0].upstream.iter().any(|reference| {
reference.table.as_deref() == Some("users") && reference.column == "id"
}));
assert_eq!(analysis.base_tables.len(), 1);
assert_eq!(analysis.base_tables[0].name, "users");
}
#[test]
fn analyze_query_reports_top_level_cte_facts() {
let analysis = analyze_query(
"WITH base(order_id, amount) AS (SELECT id, amount FROM orders), \
nested AS (WITH inner_cte AS (SELECT id FROM users) SELECT id FROM inner_cte) \
SELECT order_id FROM base",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.cte_facts.len(), 2);
assert_eq!(analysis.cte_facts[0].name, "base");
assert_eq!(analysis.cte_facts[0].columns, vec!["order_id", "amount"]);
assert!(analysis.cte_facts[0]
.body_sql
.contains("SELECT id, amount FROM orders"));
assert_eq!(analysis.cte_facts[0].output_columns, vec!["id", "amount"]);
assert_eq!(analysis.cte_facts[1].name, "nested");
assert!(analysis.cte_facts.iter().all(|cte| cte.name != "inner_cte"));
}
#[test]
fn analyze_query_reports_original_cte_body_sql_before_schema_rewrites() {
let analysis = analyze_query(
"WITH base AS (SELECT amount FROM orders) SELECT amount FROM base",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.cte_facts.len(), 1);
assert_eq!(analysis.cte_facts[0].body_sql, "SELECT amount FROM orders");
}
#[test]
fn analyze_query_reports_set_operations() {
let analysis = analyze_query(
"SELECT a FROM x UNION ALL SELECT b FROM y",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.shape, QueryShape::SetOperation);
assert_eq!(analysis.set_operations.len(), 1);
assert_eq!(analysis.set_operations[0].kind, "union");
assert!(analysis.set_operations[0].all);
assert_eq!(analysis.set_operations[0].output_columns, vec!["a"]);
assert_eq!(analysis.set_operations[0].branches.len(), 2);
assert_eq!(
analysis.set_operations[0]
.branches
.iter()
.map(|branch| branch.role)
.collect::<Vec<_>>(),
vec![SetOperationBranchRole::Value, SetOperationBranchRole::Value]
);
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "x"));
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "y"));
assert!(analysis
.base_tables
.iter()
.any(|relation| relation.name == "x"));
assert!(analysis
.base_tables
.iter()
.any(|relation| relation.name == "y"));
assert_eq!(
analysis.set_operations[0].branches[0].projections[0]
.name
.as_deref(),
Some("a")
);
for operator in ["EXCEPT", "INTERSECT"] {
let analysis = analyze_query(
&format!("SELECT a FROM x {operator} SELECT b FROM y"),
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(
analysis.set_operations[0]
.branches
.iter()
.map(|branch| branch.role)
.collect::<Vec<_>>(),
vec![
SetOperationBranchRole::Value,
SetOperationBranchRole::Filter
]
);
}
}
#[test]
fn analyze_query_reports_name_aligned_set_operation_outputs() {
let sql = "SELECT 1 AS left_value UNION ALL BY NAME SELECT 2 AS right_value";
for dialect in [DialectType::DuckDB, DialectType::Snowflake] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect,
schema: None,
},
)
.unwrap_or_else(|error| panic!("analysis failed for {dialect:?}: {error}"));
assert_eq!(
analysis.set_operations[0].output_columns,
vec!["left_value", "right_value"]
);
assert_eq!(
analysis
.projections
.iter()
.map(|projection| projection.name.as_deref())
.collect::<Vec<_>>(),
vec![Some("left_value"), Some("right_value")]
);
assert!(analysis
.projections
.iter()
.all(|projection| projection.nullability == ProjectionNullability::Nullable));
}
}
#[test]
fn analyze_query_reports_bigquery_name_alignment_modes() {
for (sql, expected_names, expected_nullability) in [
(
"SELECT 1 AS a, 2 AS b UNION ALL BY NAME SELECT 3 AS b, 4 AS a",
vec!["a", "b"],
vec![
ProjectionNullability::NonNull,
ProjectionNullability::NonNull,
],
),
(
"SELECT 1 AS a, 2 AS b INNER UNION ALL BY NAME SELECT 3 AS b, 4 AS c",
vec!["b"],
vec![ProjectionNullability::NonNull],
),
(
"SELECT 1 AS a, 2 AS b FULL OUTER UNION ALL BY NAME SELECT 3 AS b, 4 AS c",
vec!["a", "b", "c"],
vec![
ProjectionNullability::Nullable,
ProjectionNullability::NonNull,
ProjectionNullability::Nullable,
],
),
(
"SELECT 1 AS a, 2 AS b FULL OUTER UNION ALL BY NAME ON (c, a) SELECT 3 AS b, 4 AS c",
vec!["c", "a"],
vec![
ProjectionNullability::Nullable,
ProjectionNullability::Nullable,
],
),
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::BigQuery,
schema: None,
},
)
.unwrap_or_else(|error| panic!("analysis failed for {sql:?}: {error}"));
assert_eq!(analysis.set_operations[0].output_columns, expected_names);
assert_eq!(
analysis
.projections
.iter()
.map(|projection| projection.nullability)
.collect::<Vec<_>>(),
expected_nullability
);
}
}
#[test]
fn analyze_query_rejects_non_query_statements() {
let err = analyze_query(
"CREATE TABLE t (a INT)",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: None,
},
)
.unwrap_err();
assert!(err.to_string().contains("requires a SELECT"));
}
#[test]
fn analyze_query_preserves_physical_table_aliases_in_lineage() {
let analysis = analyze_query(
"SELECT o.id FROM orders AS o",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
let reference = analysis.projections[0]
.upstream
.iter()
.find(|reference| reference.column == "id")
.unwrap();
assert_eq!(reference.source_name.as_deref(), Some("orders"));
assert_eq!(reference.source_alias.as_deref(), Some("o"));
assert_eq!(reference.table.as_deref(), Some("orders"));
}
#[test]
fn analyze_query_limits_qualified_star_to_matching_source() {
let analysis = analyze_query(
"SELECT o.* FROM orders AS o JOIN customers AS c ON o.customer_id = c.id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 5);
let mut projection_names: Vec<_> = analysis
.projections
.iter()
.filter_map(|projection| projection.name.as_deref())
.collect();
projection_names.sort_unstable();
assert_eq!(
projection_names,
vec!["amount", "customer_id", "id", "total", "user_id"]
);
assert!(analysis
.projections
.iter()
.all(|projection| !projection.is_star));
assert_eq!(analysis.star_projections.len(), 1);
assert_eq!(analysis.star_projections[0].index, 0);
assert_eq!(analysis.star_projections[0].table.as_deref(), Some("o"));
let mut expanded = analysis.star_projections[0].expanded_columns.clone();
expanded.sort();
assert_eq!(
expanded,
vec!["amount", "customer_id", "id", "total", "user_id"]
);
assert!(analysis.projections.iter().all(|projection| {
projection
.upstream
.iter()
.all(|reference| reference.table.as_deref() == Some("orders"))
}));
}
#[test]
fn analyze_query_resolves_unique_unqualified_columns_with_alias_schema() {
let analysis = analyze_query(
"SELECT amount FROM orders AS o JOIN customers AS c ON o.customer_id = c.id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
let reference = analysis.projections[0]
.upstream
.iter()
.find(|reference| reference.column == "amount")
.unwrap();
assert_eq!(reference.table.as_deref(), Some("orders"));
assert_eq!(reference.source_alias.as_deref(), Some("o"));
assert_eq!(reference.confidence, ReferenceConfidence::Resolved);
}
#[test]
fn analyze_query_resolves_natural_join_merged_column_with_schema() {
let natural_join_schema: ValidationSchema = serde_json::from_value(json!({
"tables": [
{
"name": "source_table",
"columns": [
{"name": "shared_key", "type": "VARCHAR"}
]
}
]
}))
.unwrap();
let analysis = analyze_query(
"SELECT shared_key AS output_key FROM source_table \
NATURAL JOIN (SELECT shared_key FROM source_table) derived",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(natural_join_schema),
},
)
.unwrap();
let projection = &analysis.projections[0];
assert_eq!(projection.name.as_deref(), Some("output_key"));
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"));
assert!(
!projection.upstream.is_empty(),
"merged column should retain physical upstreams"
);
assert!(projection.upstream.iter().all(|reference| {
reference.column == "shared_key"
&& reference.source_name.as_deref() == Some("source_table")
&& reference.table.as_deref() == Some("source_table")
&& reference.source_kind == SourceKind::Table
&& !reference.unqualified
&& reference.confidence == ReferenceConfidence::Resolved
}));
}
#[test]
fn analyze_query_propagates_schema_type_through_anonymous_derived_table() {
let derived_table_schema: ValidationSchema = serde_json::from_value(json!({
"tables": [
{
"name": "source_table",
"columns": [
{"name": "source_value", "type": "VARCHAR"}
]
}
]
}))
.unwrap();
for sql in [
"SELECT source_value FROM (SELECT source_value FROM source_table)",
"SELECT derived.source_value FROM (SELECT source_value FROM source_table) AS derived",
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(derived_table_schema.clone()),
},
)
.unwrap();
let projection = &analysis.projections[0];
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"), "{sql}");
assert!(
projection.upstream.iter().any(|reference| {
reference.column == "source_value"
&& reference.source_name.as_deref() == Some("source_table")
&& reference.table.as_deref() == Some("source_table")
&& reference.source_kind == SourceKind::Table
&& reference.confidence == ReferenceConfidence::Resolved
}),
"{sql}: {:#?}",
projection.upstream
);
}
}
#[test]
fn analyze_query_preserves_precise_schema_type_hints() {
let analysis = analyze_query(
"SELECT amount FROM orders",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some("DECIMAL(10, 2)")
);
}
#[test]
fn analyze_query_expands_unqualified_star_with_schema() {
let analysis = analyze_query(
"SELECT * FROM orders",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
let mut names: Vec<_> = analysis
.projections
.iter()
.filter_map(|projection| projection.name.as_deref())
.collect();
names.sort_unstable();
assert_eq!(
names,
vec!["amount", "customer_id", "id", "total", "user_id"]
);
assert!(analysis
.projections
.iter()
.all(|projection| !projection.is_star));
assert_eq!(analysis.star_projections.len(), 1);
assert_eq!(analysis.star_projections[0].index, 0);
assert_eq!(analysis.star_projections[0].table, None);
let mut expanded = analysis.star_projections[0].expanded_columns.clone();
expanded.sort();
assert_eq!(
expanded,
vec!["amount", "customer_id", "id", "total", "user_id"]
);
}
#[test]
fn analyze_query_classifies_typed_aggregates() {
let analysis = analyze_query(
"SELECT COUNT(*) AS rows, SUM(amount) AS amount_sum FROM orders",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(
analysis.projections[0].transform_kind,
TransformKind::Aggregation
);
assert_eq!(
analysis.projections[1].transform_kind,
TransformKind::Aggregation
);
assert_eq!(
analysis.projections[0].nullability,
ProjectionNullability::NonNull
);
let duckdb_analysis = analyze_query(
"SELECT COUNT_IF(numeric_value > 0), MEDIAN(numeric_value), FIRST(numeric_value), ARG_MAX_NULL(label, numeric_value), ARG_MIN_NULL(label, numeric_value) FROM source_table",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
assert_eq!(duckdb_analysis.projections.len(), 5);
assert!(duckdb_analysis
.projections
.iter()
.all(|projection| projection.transform_kind == TransformKind::Aggregation));
}
#[test]
fn analyze_query_infers_duckdb_median_type_issue_425() {
let median_schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "values_table",
"columns": [{"name": "x", "type": "DOUBLE"}]
}]
}))
.unwrap();
let analysis = analyze_query(
"SELECT MEDIAN(x) AS median_x FROM values_table",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(median_schema),
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 1);
assert_eq!(
analysis.projections[0].transform_kind,
TransformKind::Aggregation
);
assert_eq!(analysis.projections[0].type_hint.as_deref(), Some("DOUBLE"));
}
#[test]
fn analyze_query_reports_projection_nullability() {
let analysis = analyze_query(
"SELECT \
COUNT(*) AS rows, \
1 AS one, \
NULL AS missing, \
o.amount, \
COALESCE(o.amount, 0) AS amount_fallback, \
c.name AS customer_name \
FROM orders AS o \
LEFT JOIN customers AS c ON o.customer_id = c.id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
let nullability: Vec<_> = analysis
.projections
.iter()
.map(|projection| projection.nullability)
.collect();
assert_eq!(
nullability,
vec![
ProjectionNullability::NonNull,
ProjectionNullability::NonNull,
ProjectionNullability::Nullable,
ProjectionNullability::Nullable,
ProjectionNullability::NonNull,
ProjectionNullability::Nullable,
]
);
}
#[test]
fn analyze_query_marks_outer_join_source_columns_nullable() {
let right_join = analyze_query(
"SELECT u.id FROM users AS u RIGHT JOIN orders AS o ON u.id = o.user_id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert_eq!(
right_join.projections[0].nullability,
ProjectionNullability::Nullable
);
let full_join = analyze_query(
"SELECT u.id, o.id FROM users AS u FULL JOIN orders AS o ON u.id = o.user_id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
assert!(full_join
.projections
.iter()
.all(|projection| projection.nullability == ProjectionNullability::Nullable));
}
#[test]
fn analyze_query_reports_transitive_base_tables() {
let analysis = analyze_query(
"WITH paid AS (SELECT customer_id FROM orders) \
SELECT c.name FROM customers AS c \
JOIN (SELECT customer_id FROM paid) AS p ON c.id = p.customer_id",
AnalyzeQueryOptions {
dialect: DialectType::Generic,
schema: Some(schema()),
},
)
.unwrap();
let base_table_names: Vec<_> = analysis
.base_tables
.iter()
.map(|relation| relation.name.as_str())
.collect();
assert_eq!(base_table_names, vec!["customers", "orders"]);
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "customers"));
assert!(analysis
.relations
.iter()
.any(|relation| relation.kind == SourceKind::DerivedTable));
}
#[test]
fn analyze_query_reports_structured_physical_table_identity() {
let analysis = analyze_query(
r#"SELECT id FROM "my.catalog"."my.schema"."orders.table" AS o"#,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let relation = analysis
.relations
.iter()
.find(|relation| relation.kind == SourceKind::Table)
.unwrap();
assert_eq!(relation.name, "my.catalog.my.schema.orders.table");
assert_eq!(relation.catalog.as_deref(), Some("my.catalog"));
assert_eq!(relation.schema.as_deref(), Some("my.schema"));
assert_eq!(relation.table.as_deref(), Some("orders.table"));
assert_eq!(relation.alias.as_deref(), Some("o"));
assert_eq!(analysis.base_tables.len(), 1);
assert_eq!(analysis.base_tables[0].name, relation.name);
assert_eq!(analysis.base_tables[0].catalog, relation.catalog);
assert_eq!(analysis.base_tables[0].schema, relation.schema);
assert_eq!(analysis.base_tables[0].table, relation.table);
}
#[test]
fn analyze_query_reports_structured_table_identity_for_qualified_and_derived_sources() {
let analysis = analyze_query(
"SELECT x FROM (SELECT id AS x FROM mycatalog.myschema.orders) d",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let derived = analysis
.relations
.iter()
.find(|relation| relation.name == "d")
.unwrap();
assert_eq!(derived.kind, SourceKind::DerivedTable);
assert_eq!(derived.catalog, None);
assert_eq!(derived.schema, None);
assert_eq!(derived.table, None);
assert_eq!(analysis.base_tables.len(), 1);
let base_table = &analysis.base_tables[0];
assert_eq!(base_table.name, "mycatalog.myschema.orders");
assert_eq!(base_table.catalog.as_deref(), Some("mycatalog"));
assert_eq!(base_table.schema.as_deref(), Some("myschema"));
assert_eq!(base_table.table.as_deref(), Some("orders"));
}
#[test]
fn analyze_query_reports_base_tables_inside_derived_table() {
let analysis = analyze_query(
"SELECT x FROM (SELECT id AS x FROM orders) d",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let base_table_names: Vec<_> = analysis
.base_tables
.iter()
.map(|relation| relation.name.as_str())
.collect();
assert_eq!(base_table_names, vec!["orders"]);
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "d" && relation.kind == SourceKind::DerivedTable));
}
#[test]
fn analyze_query_reports_base_tables_inside_derived_table_set_operation() {
let analysis = analyze_query(
"SELECT s FROM (SELECT a AS s FROM orders UNION ALL SELECT a AS s FROM users) u",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let base_table_names: Vec<_> = analysis
.base_tables
.iter()
.map(|relation| relation.name.as_str())
.collect();
assert_eq!(base_table_names, vec!["orders", "users"]);
assert!(analysis
.relations
.iter()
.any(|relation| relation.name == "u" && relation.kind == SourceKind::DerivedTable));
}
fn single_column_schema(table_names: &[&str], column_name: &str) -> ValidationSchema {
let tables: Vec<_> = table_names
.iter()
.map(|name| {
json!({
"name": name,
"columns": [{"name": column_name, "type": "INT"}]
})
})
.collect();
serde_json::from_value(json!({ "tables": tables })).unwrap()
}
fn unnest_analysis_schema() -> ValidationSchema {
serde_json::from_value(json!({
"tables": [
{
"name": "t",
"columns": [{"name": "arr", "type": "INT"}]
}
]
}))
.unwrap()
}
fn struct_field_analysis_schema() -> ValidationSchema {
serde_json::from_value(json!({
"tables": [{
"name": "source_table",
"columns": [
{"name": "nested_items", "type": "STRUCT(field_value VARCHAR)[]"},
{"name": "composite_value", "type": "STRUCT(field_value VARCHAR, label VARCHAR)"}
]
}]
}))
.unwrap()
}
#[test]
fn analyze_query_resolves_nested_set_operation_inside_derived_table() {
let analysis = analyze_query(
"SELECT v FROM ((SELECT v FROM t1 UNION ALL SELECT v FROM t2) \
UNION ALL SELECT v FROM t3) u",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(single_column_schema(&["t1", "t2", "t3"], "v")),
},
)
.unwrap();
let upstream_tables: Vec<_> = analysis.projections[0]
.upstream
.iter()
.filter_map(|reference| reference.table.as_deref())
.collect();
assert!(upstream_tables.contains(&"t1"));
assert!(upstream_tables.contains(&"t2"));
assert!(upstream_tables.contains(&"t3"));
let base_table_names: Vec<_> = analysis
.base_tables
.iter()
.map(|relation| relation.name.as_str())
.collect();
assert_eq!(base_table_names, vec!["t1", "t2", "t3"]);
}
#[test]
fn analyze_query_resolves_mixed_nested_set_operation_arm_inside_derived_table() {
let analysis = analyze_query(
"SELECT v FROM (SELECT v FROM t0 UNION ALL \
(SELECT v FROM t1 UNION ALL SELECT v FROM t2)) u",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(single_column_schema(&["t0", "t1", "t2"], "v")),
},
)
.unwrap();
let upstream_tables: Vec<_> = analysis.projections[0]
.upstream
.iter()
.filter_map(|reference| reference.table.as_deref())
.collect();
assert!(upstream_tables.contains(&"t0"));
assert!(upstream_tables.contains(&"t1"));
assert!(upstream_tables.contains(&"t2"));
}
#[test]
fn analyze_query_resolves_nested_set_operation_inside_cte_with_schema() {
let analysis = analyze_query(
"WITH c AS (SELECT v FROM ((SELECT v FROM t1 UNION ALL SELECT v FROM t2) \
UNION ALL SELECT v FROM t3) u) SELECT v FROM c",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(single_column_schema(&["t1", "t2", "t3"], "v")),
},
)
.unwrap();
let upstream_tables: Vec<_> = analysis.projections[0]
.upstream
.iter()
.filter_map(|reference| reference.table.as_deref())
.collect();
assert!(upstream_tables.contains(&"t1"));
assert!(upstream_tables.contains(&"t2"));
assert!(upstream_tables.contains(&"t3"));
}
#[test]
fn analyze_query_resolves_unnest_virtual_output_aliases_with_schema() {
for sql in [
"SELECT i FROM t, UNNEST(t.arr) AS i",
"SELECT i FROM t, UNNEST(t.arr) AS u(i)",
"SELECT u.i FROM t, UNNEST(t.arr) AS u(i)",
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(unnest_analysis_schema()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert!(
analysis.projections[0]
.upstream
.iter()
.any(|reference| reference.table.as_deref() == Some("t")
&& reference.column == "arr"),
"expected t.arr upstream for {sql:?}, got {:?}",
analysis.projections[0].upstream
);
}
}
#[test]
fn analyze_query_resolves_struct_fields_and_types_issue_408() {
let cases = [
(
"SELECT composite_value.field_value AS output_value FROM source_table",
"composite_value",
),
(
"SELECT source_table.composite_value.field_value AS output_value FROM source_table",
"composite_value",
),
(
"SELECT item.field_value AS output_value FROM source_table s \
CROSS JOIN UNNEST(s.nested_items) AS expanded(item)",
"nested_items",
),
];
for (sql, expected_column) in cases {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(struct_field_analysis_schema()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
let projection = &analysis.projections[0];
assert_eq!(projection.type_hint.as_deref(), Some("TEXT"));
assert!(
projection.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("source_table")
&& reference.column == expected_column
}),
"expected source_table.{expected_column} upstream for {sql:?}, got {:?}",
projection.upstream
);
}
}
#[test]
fn analyze_query_propagates_unnest_element_types_through_query_scopes() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "events",
"columns": [{"name": "tags", "type": "VARCHAR[]"}]
}]
}))
.unwrap();
for sql in [
"SELECT UNNEST(e.tags) AS tag FROM events AS e",
"SELECT u.tag FROM events AS e, UNNEST(e.tags) AS u(tag)",
"WITH exploded AS (SELECT u.tag AS tag FROM events AS e, \
UNNEST(e.tags) AS u(tag)) SELECT tag FROM exploded",
"SELECT d.tag FROM (SELECT u.tag AS tag FROM events AS e, \
UNNEST(e.tags) AS u(tag)) AS d",
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema.clone()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some("TEXT"),
"unexpected output type for {sql:?}"
);
}
}
#[test]
fn analyze_query_infers_unnest_type_from_case_array_constructor() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "events",
"columns": [
{"name": "created_at", "type": "TIMESTAMP"},
{"name": "closed_at", "type": "TIMESTAMP"}
]
}]
}))
.unwrap();
let analysis = analyze_query(
"SELECT UNNEST(CASE WHEN closed_at IS NULL THEN ARRAY[created_at] \
ELSE ARRAY[created_at, closed_at] END) AS event_at FROM events",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema),
},
)
.unwrap();
assert_eq!(analysis.projections[0].name.as_deref(), Some("event_at"));
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some("TIMESTAMP")
);
assert!(analysis.projections[0].upstream.iter().any(|reference| {
reference.table.as_deref() == Some("events") && reference.column == "created_at"
}));
assert!(analysis.projections[0].upstream.iter().any(|reference| {
reference.table.as_deref() == Some("events") && reference.column == "closed_at"
}));
}
#[test]
fn analyze_query_preserves_type_through_parenthesized_case() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [{
"name": "flags",
"columns": [{"name": "flag", "type": "BOOLEAN"}]
}]
}))
.unwrap();
for sql in [
"SELECT CASE WHEN flag THEN 'yes' ELSE 'no' END AS label FROM flags",
"SELECT (CASE WHEN flag THEN 'yes' ELSE 'no' END) AS label FROM flags",
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema.clone()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert_eq!(analysis.projections[0].name.as_deref(), Some("label"));
assert_eq!(
analysis.projections[0].type_hint.as_deref(),
Some("TEXT"),
"unexpected output type for {sql:?}"
);
assert!(analysis.projections[0].upstream.iter().any(|reference| {
reference.table.as_deref() == Some("flags") && reference.column == "flag"
}));
}
}
#[test]
fn analyze_query_keeps_reused_unnest_alias_types_isolated_between_ctes() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [
{
"name": "integer_events",
"columns": [{"name": "items", "type": "INT[]"}]
},
{
"name": "text_events",
"columns": [{"name": "items", "type": "VARCHAR[]"}]
}
]
}))
.unwrap();
for sql in [
"WITH ints AS (SELECT u.item FROM integer_events AS e, \
UNNEST(e.items) AS u(item)), \
texts AS (SELECT u.item FROM text_events AS e, \
UNNEST(e.items) AS u(item)) \
SELECT ints.item AS int_item, texts.item AS text_item FROM ints CROSS JOIN texts",
"WITH texts AS (SELECT u.item FROM text_events AS e, \
UNNEST(e.items) AS u(item)), \
ints AS (SELECT u.item FROM integer_events AS e, \
UNNEST(e.items) AS u(item)) \
SELECT ints.item AS int_item, texts.item AS text_item FROM ints CROSS JOIN texts",
] {
let analysis = analyze_query(
sql,
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema.clone()),
},
)
.unwrap_or_else(|error| panic!("analyze_query failed for {sql:?}: {error}"));
assert_eq!(analysis.projections[0].type_hint.as_deref(), Some("INT"));
assert_eq!(analysis.projections[1].type_hint.as_deref(), Some("TEXT"));
}
}
#[test]
fn analyze_query_tolerates_partial_schema_for_unknown_columns() {
let analysis = analyze_query(
"SELECT order_id, amount FROM t",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(single_column_schema(&["t"], "amount")),
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 2);
assert_eq!(analysis.projections[0].name.as_deref(), Some("order_id"));
assert_eq!(analysis.projections[1].name.as_deref(), Some("amount"));
let order_id = &analysis.projections[0].upstream;
assert!(
order_id.iter().any(|reference| {
reference.column == "order_id"
&& reference.confidence == ReferenceConfidence::Resolved
&& reference.source_name.as_deref() == Some("t")
&& reference.table.as_deref() == Some("t")
}),
"expected best-effort order_id reference, got {order_id:?}"
);
let amount = &analysis.projections[1].upstream;
assert!(
amount.iter().any(|reference| {
reference.column == "amount"
&& reference.confidence == ReferenceConfidence::Resolved
&& reference.table.as_deref() == Some("t")
}),
"expected schema-backed amount reference, got {amount:?}"
);
}
#[test]
fn analyze_query_tolerates_partial_schema_for_qualified_unknown_columns() {
let analysis = analyze_query(
"SELECT t.order_id, t.amount FROM t",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(single_column_schema(&["t"], "amount")),
},
)
.unwrap();
assert!(
analysis.projections[0].upstream.iter().any(|reference| {
reference.column == "order_id"
&& reference.confidence == ReferenceConfidence::Resolved
&& reference.table.as_deref() == Some("t")
}),
"expected qualified unknown column to stay as best-effort t.order_id, got {:?}",
analysis.projections[0].upstream
);
assert!(
analysis.projections[1].upstream.iter().any(|reference| {
reference.column == "amount" && reference.table.as_deref() == Some("t")
}),
"expected known t.amount to resolve, got {:?}",
analysis.projections[1].upstream
);
}
#[test]
fn analyze_query_tolerates_partial_schema_for_join_conditions() {
let schema: ValidationSchema = serde_json::from_value(json!({
"tables": [
{
"name": "t",
"columns": [{"name": "order_id", "type": "INT"}]
},
{
"name": "u",
"columns": [{"name": "amount", "type": "INT"}]
}
]
}))
.unwrap();
let analysis = analyze_query(
"SELECT a.order_id, b.amount FROM t a JOIN u b ON a.id = b.id",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: Some(schema),
},
)
.unwrap();
assert_eq!(analysis.projections.len(), 2);
assert!(
analysis.projections[0].upstream.iter().any(|reference| {
reference.column == "order_id"
&& reference.source_alias.as_deref() == Some("a")
&& reference.table.as_deref() == Some("t")
}),
"expected a.order_id to resolve, got {:?}",
analysis.projections[0].upstream
);
assert!(
analysis.projections[1].upstream.iter().any(|reference| {
reference.column == "amount"
&& reference.source_alias.as_deref() == Some("b")
&& reference.table.as_deref() == Some("u")
}),
"expected b.amount to resolve, got {:?}",
analysis.projections[1].upstream
);
}
#[test]
fn analyze_query_resolves_same_select_alias_reference() {
let analysis = analyze_query(
"WITH c AS (SELECT x FROM t) SELECT c.x AS a, a + 1 AS b FROM c",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let projection = analysis
.projections
.iter()
.find(|projection| projection.name.as_deref() == Some("b"))
.unwrap();
assert!(projection
.upstream
.iter()
.any(|reference| reference.table.as_deref() == Some("t") && reference.column == "x"));
}
#[test]
fn analyze_query_resolves_pivot_alias_columns_and_generated_outputs() {
let analysis = analyze_query(
"SELECT region2, p1 FROM (SELECT region, q, amt FROM sales) \
PIVOT(SUM(amt) FOR q IN ('Q1')) AS p(region2, p1)",
AnalyzeQueryOptions {
dialect: DialectType::DuckDB,
schema: None,
},
)
.unwrap();
let region = analysis
.projections
.iter()
.find(|projection| projection.name.as_deref() == Some("region2"))
.unwrap();
assert!(region.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("sales") && reference.column == "region"
}));
let pivot_value = analysis
.projections
.iter()
.find(|projection| projection.name.as_deref() == Some("p1"))
.unwrap();
assert!(pivot_value.upstream.iter().any(|reference| {
reference.table.as_deref() == Some("sales") && reference.column == "amt"
}));
}