polyglot-sql 0.5.8__tar.gz → 0.5.9__tar.gz

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/Cargo.lock

@@ -605,7 +605,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql"
-version = "0.5.8"
+version = "0.5.9"
 dependencies = [
  "criterion",
  "once_cell",
@@ -621,7 +621,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-ffi"
-version = "0.5.8"
+version = "0.5.9"
 dependencies = [
  "cbindgen",
  "polyglot-sql",
@@ -631,11 +631,11 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-function-catalogs"
-version = "0.5.8"
+version = "0.5.9"
 
 [[package]]
 name = "polyglot-sql-python"
-version = "0.5.8"
+version = "0.5.9"
 dependencies = [
  "polyglot-sql",
  "pyo3",
@@ -646,7 +646,7 @@ dependencies = [
 
 [[package]]
 name = "polyglot-sql-wasm"
-version = "0.5.8"
+version = "0.5.9"
 dependencies = [
  "console_error_panic_hook",
  "js-sys",

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/Cargo.toml

@@ -6,7 +6,7 @@ exclude = [
 ]
 
 [workspace.package]
-version = "0.5.8"
+version = "0.5.9"
 edition = "2021"
 license = "MIT"
 authors = ["polyglot contributors"]

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polyglot-sql
-Version: 0.5.8
+Version: 0.5.9
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/crates/polyglot-sql/Cargo.toml

@@ -106,7 +106,7 @@ thiserror = { workspace = true }
 unicode-segmentation = { workspace = true }
 stacker = { version = "0.1", optional = true }
 ts-rs = { version = "12.0", features = ["serde-compat"], optional = true }
-polyglot-sql-function-catalogs = { path = "../polyglot-sql-function-catalogs", version = "0.5.8", optional = true, default-features = false }
+polyglot-sql-function-catalogs = { path = "../polyglot-sql-function-catalogs", version = "0.5.9", optional = true, default-features = false }
 
 [dev-dependencies]
 pretty_assertions = "1.4"

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/crates/polyglot-sql/src/resolver.rs

@@ -173,7 +173,14 @@ impl<'a> Resolver<'a> {
 
     /// Extract column names from a source expression
     fn extract_columns_from_source(&self, source_info: &SourceInfo) -> ResolverResult<Vec<String>> {
-        let columns = match &source_info.expression {
+        self.get_source_columns_for_expression(&source_info.expression)
+    }
+
+    fn get_source_columns_for_expression(
+        &self,
+        expression: &Expression,
+    ) -> ResolverResult<Vec<String>> {
+        let columns = match expression {
             Expression::Table(table) => {
                 // For tables, try to get columns from schema.
                 // Build the fully qualified name (catalog.schema.table) to
@@ -211,6 +218,17 @@ impl<'a> Resolver<'a> {
             }
             Expression::Pivot(pivot) => self.get_pivot_output_columns(pivot),
             Expression::Unpivot(unpivot) => self.get_unpivot_output_columns(unpivot),
+            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
+                alias_output_columns(alias)
+            }
+            Expression::Alias(alias) => {
+                let columns = self.get_source_columns_for_expression(&alias.this)?;
+                apply_alias_columns(columns, &alias.column_aliases)
+            }
+            Expression::Unnest(unnest) => unnest_output_columns(unnest),
+            Expression::Lateral(lateral) => lateral_output_columns(lateral),
+            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
+            Expression::Paren(paren) => self.get_source_columns_for_expression(&paren.this)?,
             _ => Vec::new(),
         };
 
@@ -228,6 +246,11 @@ impl<'a> Resolver<'a> {
             Expression::Intersect(intersect) => self.get_named_selects(&intersect.left),
             Expression::Except(except) => self.get_named_selects(&except.left),
             Expression::Subquery(subquery) => self.get_named_selects(&subquery.this),
+            Expression::Alias(alias) => {
+                let columns = self.get_named_selects(&alias.this);
+                apply_alias_columns(columns, &alias.column_aliases)
+            }
+            Expression::Paren(paren) => self.get_named_selects(&paren.this),
             _ => Vec::new(),
         }
     }
@@ -351,6 +374,16 @@ impl<'a> Resolver<'a> {
             Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) => self
                 .get_source_columns_from_set_op(source)
                 .unwrap_or_default(),
+            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
+                alias_output_columns(alias)
+            }
+            Expression::Alias(alias) => {
+                let columns = self.get_source_output_columns(&alias.this);
+                apply_alias_columns(columns, &alias.column_aliases)
+            }
+            Expression::Unnest(unnest) => unnest_output_columns(unnest),
+            Expression::Lateral(lateral) => lateral_output_columns(lateral),
+            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
             Expression::Cte(cte) => {
                 if cte.columns.is_empty() {
                     self.get_named_selects(&cte.this)
@@ -383,6 +416,11 @@ impl<'a> Resolver<'a> {
                     Ok(self.get_named_selects(&subquery.this))
                 }
             }
+            Expression::Alias(alias) => {
+                let columns = self.get_source_columns_from_set_op(&alias.this)?;
+                Ok(apply_alias_columns(columns, &alias.column_aliases))
+            }
+            Expression::Paren(paren) => self.get_source_columns_from_set_op(&paren.this),
             Expression::Union(union) => {
                 // Standard UNION: columns come from the left side
                 self.get_source_columns_from_set_op(&union.left)
@@ -540,6 +578,59 @@ fn apply_alias_columns(mut columns: Vec<String>, alias_columns: &[Identifier]) -
     columns
 }
 
+fn unnest_output_columns(unnest: &crate::expressions::UnnestFunc) -> Vec<String> {
+    unnest
+        .alias
+        .iter()
+        .map(|alias| alias.name.clone())
+        .chain(unnest.offset_alias.iter().map(|alias| alias.name.clone()))
+        .collect()
+}
+
+fn alias_output_columns(alias: &crate::expressions::Alias) -> Vec<String> {
+    if alias.column_aliases.is_empty() {
+        vec![alias.alias.name.clone()]
+    } else {
+        alias
+            .column_aliases
+            .iter()
+            .map(|column| column.name.clone())
+            .collect()
+    }
+}
+
+fn lateral_output_columns(lateral: &crate::expressions::Lateral) -> Vec<String> {
+    if lateral.column_aliases.is_empty() {
+        default_virtual_output_columns(&lateral.this)
+    } else {
+        lateral.column_aliases.clone()
+    }
+}
+
+fn lateral_view_output_columns(lateral_view: &crate::expressions::LateralView) -> Vec<String> {
+    lateral_view
+        .column_aliases
+        .iter()
+        .map(|column| column.name.clone())
+        .collect()
+}
+
+fn default_virtual_output_columns(expression: &Expression) -> Vec<String> {
+    match expression {
+        Expression::Unnest(unnest) => unnest_output_columns(unnest),
+        Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
+            alias_output_columns(alias)
+        }
+        Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN") => {
+            ["seq", "key", "path", "index", "value", "this"]
+                .into_iter()
+                .map(String::from)
+                .collect()
+        }
+        _ => Vec::new(),
+    }
+}
+
 fn pivot_excluded_source_columns(
     pivot: &crate::expressions::Pivot,
     dialect: Option<DialectType>,

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/crates/polyglot-sql/tests/query_analysis.rs

@@ -609,6 +609,32 @@ fn analyze_query_reports_base_tables_inside_derived_table_set_operation() {
         .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()
+}
+
 #[test]
 fn analyze_query_resolves_nested_set_operation_inside_derived_table() {
     let analysis = analyze_query(
@@ -616,7 +642,7 @@ fn analyze_query_resolves_nested_set_operation_inside_derived_table() {
          UNION ALL SELECT v FROM t3) u",
         AnalyzeQueryOptions {
             dialect: DialectType::DuckDB,
-            schema: None,
+            schema: Some(single_column_schema(&["t1", "t2", "t3"], "v")),
         },
     )
     .unwrap();
@@ -638,6 +664,78 @@ fn analyze_query_resolves_nested_set_operation_inside_derived_table() {
     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_same_select_alias_reference() {
     let analysis = analyze_query(

{polyglot_sql-0.5.8 → polyglot_sql-0.5.9}/crates/polyglot-sql-python/tests/test_query_analysis.py

@@ -148,6 +148,46 @@ def test_analyze_query_resolves_pivot_alias_columns():
     )
 
 
+def test_analyze_query_resolves_nested_set_operation_with_schema():
+    result = polyglot_sql.analyze_query(
+        "SELECT v FROM ((SELECT v FROM t1 UNION ALL SELECT v FROM t2) "
+        "UNION ALL SELECT v FROM t3) u",
+        {
+            "dialect": "duckdb",
+            "schema": {
+                "tables": [
+                    {"name": "t1", "columns": [{"name": "v", "type": "INT"}]},
+                    {"name": "t2", "columns": [{"name": "v", "type": "INT"}]},
+                    {"name": "t3", "columns": [{"name": "v", "type": "INT"}]},
+                ]
+            },
+        },
+    )
+
+    upstream = result["projections"][0]["upstream"]
+    assert {reference["table"] for reference in upstream} == {"t1", "t2", "t3"}
+
+
+def test_analyze_query_resolves_unnest_output_alias_with_schema():
+    result = polyglot_sql.analyze_query(
+        "SELECT i FROM t, UNNEST(t.arr) AS i",
+        {
+            "dialect": "duckdb",
+            "schema": {
+                "tables": [
+                    {"name": "t", "columns": [{"name": "arr", "type": "INT"}]},
+                ]
+            },
+        },
+    )
+
+    upstream = result["projections"][0]["upstream"]
+    assert any(
+        reference["table"] == "t" and reference["column"] == "arr"
+        for reference in upstream
+    )
+
+
 def test_analyze_query_unknown_dialect_raises_value_error():
     with pytest.raises(ValueError):
         polyglot_sql.analyze_query("SELECT 1", dialect="not_a_dialect")