feat(native): recover CTE facts natively behind the preview engine

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs

@@ -1,10 +1,14 @@
-//! Python's dict walk over a parsed query's `to_dict()` payload, on the same serde values.
+//! Python's view of a parsed query: the wheel's expression accessors and its dict walk.
 
-use serde_json::{Map, Value};
+use polyglot_sql::expressions::Cte;
+use polyglot_sql::traversal::ExpressionWalk;
+use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions, ast_json};
+use serde_json::{Map, Value, json};
 
 use crate::assembly::analysis_session::constants::{
-    COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, FUNCTION_AST_KIND, RENDERED_TYPE_NAMES,
-    STAR_AST_KIND, TABLE_AST_KIND,
+    ALIAS_AST_KIND, ANNOTATED_AST_KIND, COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES,
+    FUNCTION_AST_KIND, MAX_FUNCTION_CALL_DEPTH, RENDERED_TYPE_NAMES, STAR_AST_KIND, TABLE_AST_KIND,
+    WILDCARD,
 };
 
 /// Python's `bool(value)`.
@@ -142,7 +146,7 @@ pub(crate) fn render_type(node: Option<&Value>) -> Option<String> {
 }
 
 /// Python's `str(value)` for a value `isinstance(value, int)` accepts.
-fn python_int(value: Option<&Value>) -> Option<String> {
+pub(crate) fn python_int(value: Option<&Value>) -> Option<String> {
     match value? {
         Value::Bool(flag) => Some(if *flag { "True" } else { "False" }.to_owned()),
         Value::Number(number) if number.is_i64() || number.is_u64() => Some(number.to_string()),
@@ -154,3 +158,236 @@ fn python_int(value: Option<&Value>) -> Option<String> {
 pub(crate) fn casefold(text: &str) -> String {
     text.to_ascii_lowercase()
 }
+
+/// A value the wheel hands Python for one payload key: an expression, a list of them, or data.
+#[derive(Debug, Clone)]
+pub(crate) enum PyValue {
+    Expr(Expression),
+    Exprs(Vec<Expression>),
+    Raw(Value),
+}
+
+/// The wheel's guarded `parse_one`; Ok(None) is the `PolyglotError` Python catches.
+pub(crate) fn parse_one(sql: &str, dialect: &str) -> Result<Option<Expression>, String> {
+    let guard: ComplexityGuardOptions =
+        serde_json::from_value(json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH}))
+            .map_err(|error| error.to_string())?;
+    let options: ParseOptions = ParseOptions {
+        complexity_guard: Some(guard),
+    };
+    let parsed: Result<Vec<Expression>, _> = Dialect::get_by_name(dialect)
+        .ok_or_else(|| format!("unknown dialect {dialect}"))?
+        .parse_with_options(sql, &options);
+    let mut expressions: Vec<Expression> = match parsed {
+        Ok(expressions) => expressions,
+        Err(_) => return Ok(None),
+    };
+    if expressions.len() != 1 {
+        return Ok(None);
+    }
+    Ok(expressions.pop())
+}
+
+/// `str(getattr(node, "kind", ""))`.
+pub(crate) fn kind(node: Option<&Expression>) -> &'static str {
+    node.map_or("", Expression::variant_name)
+}
+
+/// `node.name`, "" for None.
+pub(crate) fn name(node: Option<&Expression>) -> &str {
+    node.map_or("", Expression::get_name)
+}
+
+/// `node.output_name`, "" for None.
+pub(crate) fn output_name(node: Option<&Expression>) -> &str {
+    node.map_or("", Expression::get_output_name)
+}
+
+/// `node.alias_or_name`.
+pub(crate) fn alias_or_name(node: &Expression) -> &str {
+    let alias: &str = node.get_alias();
+    if alias.is_empty() {
+        node.get_name()
+    } else {
+        alias
+    }
+}
+
+/// `node.is_star`.
+pub(crate) fn is_star(node: Option<&Expression>) -> bool {
+    match node {
+        Some(Expression::Star(_)) => true,
+        Some(Expression::Column(column)) => column.name.name == WILDCARD,
+        _ => false,
+    }
+}
+
+/// Python's `_unwrap_polyglot_annotations`.
+pub(crate) fn unwrap_annotations(mut node: Option<&Expression>) -> Option<&Expression> {
+    while kind(node) == ANNOTATED_AST_KIND {
+        node = node.and_then(Expression::get_this);
+    }
+    node
+}
+
+/// A projection's expression: an alias's `this`, otherwise the projection.
+pub(crate) fn projected_expression(projection: Option<&Expression>) -> Option<&Expression> {
+    if kind(projection) == ALIAS_AST_KIND {
+        return projection.and_then(Expression::get_this);
+    }
+    projection
+}
+
+/// `node.this` then `node.expressions`: Python's `_polyglot_expression_args`.
+pub(crate) fn expression_args(node: Option<&Expression>) -> Vec<&Expression> {
+    let Some(node) = node else {
+        return Vec::new();
+    };
+    node.get_this()
+        .into_iter()
+        .chain(node.get_expressions())
+        .collect()
+}
+
+/// `node.children()`.
+pub(crate) fn children(node: &Expression) -> Vec<&Expression> {
+    ExpressionWalk::children(node)
+}
+
+/// Python's `_polyglot_direct_select_tables`: the table children.
+pub(crate) fn direct_tables(node: &Expression) -> Vec<&Expression> {
+    children(node)
+        .into_iter()
+        .filter(|child| child.variant_name() == TABLE_AST_KIND)
+        .collect()
+}
+
+/// `node.find_all(kind)`: descendants in depth-first order, the node itself excluded.
+pub(crate) fn find_all<'a>(node: &'a Expression, wanted: &str) -> Vec<&'a Expression> {
+    node.dfs()
+        .skip(1)
+        .filter(|descendant| descendant.variant_name() == wanted)
+        .collect()
+}
+
+/// `node.to_dict()`.
+pub(crate) fn to_dict(node: &Expression) -> Result<Value, String> {
+    serde_json::to_value(node).map_err(|error| error.to_string())
+}
+
+/// The wheel's `expression_payload`: the tagged value's fields.
+pub(crate) fn expression_payload(node: &Expression) -> Result<Map<String, Value>, String> {
+    Ok(match to_dict(node)? {
+        Value::Object(map) => match map.into_iter().next() {
+            Some((_, Value::Object(payload))) => payload,
+            _ => Map::new(),
+        },
+        _ => Map::new(),
+    })
+}
+
+/// `node.args.get(key)` (and `node.arg(key)` outside its special keys); None for Python's None.
+pub(crate) fn arg(node: Option<&Expression>, key: &str) -> Result<Option<PyValue>, String> {
+    let Some(node) = node else {
+        return Ok(None);
+    };
+    Ok(expression_payload(node)?.remove(key).and_then(py_value))
+}
+
+/// The wheel's `value_to_python_object`; None where Python sees None.
+fn py_value(value: Value) -> Option<PyValue> {
+    if let Ok(expression) = ast_json::expression_from_value(value.clone()) {
+        return Some(PyValue::Expr(expression));
+    }
+    if let Ok(expressions) = ast_json::expressions_from_value(value.clone()) {
+        return Some(PyValue::Exprs(expressions));
+    }
+    (!value.is_null()).then_some(PyValue::Raw(value))
+}
+
+/// Python's `bool()` of a converted value: expressions are always truthy.
+pub(crate) fn py_truthy(value: Option<&PyValue>) -> bool {
+    match value {
+        None => false,
+        Some(PyValue::Expr(_)) => true,
+        Some(PyValue::Exprs(items)) => !items.is_empty(),
+        Some(PyValue::Raw(raw)) => truthy(Some(raw)),
+    }
+}
+
+/// Python's `_polyglot_name_payload_value` of a converted value.
+pub(crate) fn name_payload(value: Option<&PyValue>) -> String {
+    match value {
+        Some(PyValue::Raw(raw)) => raw_name_payload(Some(raw)),
+        _ => String::new(),
+    }
+}
+
+/// Python's `_polyglot_name_payload_value` of plain data.
+pub(crate) fn raw_name_payload(value: Option<&Value>) -> String {
+    match value {
+        Some(Value::String(text)) => text.clone(),
+        Some(Value::Object(object)) => object
+            .get("name")
+            .and_then(Value::as_str)
+            .unwrap_or_default()
+            .to_owned(),
+        _ => String::new(),
+    }
+}
+
+/// Python's `_polyglot_column_table_name`: `to_dict()["column"]["table"]["name"]`.
+pub(crate) fn column_table_name(node: &Expression) -> Result<String, String> {
+    let value: Value = to_dict(node)?;
+    let table: Option<&Value> = value
+        .get("column")
+        .filter(|column| column.is_object())
+        .and_then(|column| column.get("table"))
+        .filter(|table| table.is_object());
+    Ok(raw_name_payload(table))
+}
+
+/// The top-level CTEs `with_ctes()` returns: name, whether it has column aliases, and its body.
+pub(crate) fn top_level_ctes(node: &Expression) -> Vec<(&str, bool, &Expression)> {
+    let with = match node {
+        Expression::Select(select) => select.with.as_ref(),
+        Expression::Union(union) => union.with.as_ref(),
+        Expression::Intersect(intersect) => intersect.with.as_ref(),
+        Expression::Except(except) => except.with.as_ref(),
+        _ => None,
+    };
+    with.map(|with| with.ctes.iter().map(cte_entry).collect())
+        .unwrap_or_default()
+}
+
+fn cte_entry(cte: &Cte) -> (&str, bool, &Expression) {
+    (cte.alias.name.as_str(), !cte.columns.is_empty(), &cte.this)
+}
+
+/// A set operation's `arg("left")`, `arg("right")` and `arg("by_name")`.
+pub(crate) fn set_operation(node: &Expression) -> Option<(&Expression, &Expression, bool)> {
+    match node {
+        Expression::Union(operation) => {
+            Some((&operation.left, &operation.right, operation.by_name))
+        }
+        Expression::Intersect(operation) => {
+            Some((&operation.left, &operation.right, operation.by_name))
+        }
+        Expression::Except(operation) => {
+            Some((&operation.left, &operation.right, operation.by_name))
+        }
+        _ => None,
+    }
+}
+
+/// A select's `arg("where_clause")` as plain data.
+pub(crate) fn where_clause(node: &Expression) -> Result<Option<Value>, String> {
+    let Expression::Select(select) = node else {
+        return Ok(None);
+    };
+    select
+        .where_clause
+        .as_ref()
+        .map(|clause| serde_json::to_value(clause).map_err(|error| error.to_string()))
+        .transpose()
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs

@@ -2,17 +2,16 @@
 
 use std::collections::{HashMap, HashSet};
 
-use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions};
-use serde_json::{Map, Value, json};
+use serde_json::{Map, Value};
 use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic;
 
 use crate::assembly::analysis_session::_helpers::dict_walk::{
     casefold, column_name, dict_list, identifier_name, is_single_wildcard, nested, node_key,
-    payload, relation_name, render_type, truthy,
+    parse_one, payload, relation_name, render_type, to_dict, truthy,
 };
 use crate::assembly::analysis_session::constants::{
     ALIAS_AST_KIND, CAST_AST_KIND, CTE_AST_KIND, DUCKDB_DIALECT, DYNAMIC_VALUE_SOURCE_KINDS,
-    MAX_FUNCTION_CALL_DEPTH, MOTHERDUCK_DIALECT, PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND,
+    MOTHERDUCK_DIALECT, PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND,
     SIMPLIFIED_PIVOT_DIALECTS, SUPPORTED_PIVOT_DIALECTS, TYPE_PRESERVING_AGGREGATES,
     UNKNOWN_NULLABILITY, UNKNOWN_TYPE,
 };
@@ -80,25 +79,9 @@ pub(crate) fn pivot_outcome(
 
 /// Python's `parse_one(...).to_dict()`, or None where Python reports a parse error.
 fn parsed(sql: &str, dialect: &str) -> Result<Option<Value>, String> {
-    let guard: ComplexityGuardOptions =
-        serde_json::from_value(json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH}))
-            .map_err(|error| error.to_string())?;
-    let options: ParseOptions = ParseOptions {
-        complexity_guard: Some(guard),
-    };
-    let parsed: Result<Vec<Expression>, _> = Dialect::get_by_name(dialect)
-        .ok_or_else(|| format!("unknown dialect {dialect}"))?
-        .parse_with_options(sql, &options);
-    let expressions: Vec<Expression> = match parsed {
-        Ok(expressions) => expressions,
-        Err(_) => return Ok(None),
-    };
-    let [expression] = expressions.as_slice() else {
-        return Ok(None);
-    };
-    serde_json::to_value(expression)
-        .map(Some)
-        .map_err(|error| error.to_string())
+    parse_one(sql, dialect)?
+        .map(|expression| to_dict(&expression))
+        .transpose()
 }
 
 /// Casefolding equals Python's only for ASCII text.