[Traversal] transform_map does not visit the children of most typed function nodes

[aniket-wisdom]

Rewrites that go through transform_map never see the columns inside most typed function nodes. rename_columns renames the column in UPPER(t.c) but leaves INSTR(t.c, 'x'), NVL2(t.c, 1, 0), LAST_DAY(t.c), YEAR(t.c) and STARTS_WITH(t.c, 'a') as they are, and nothing signals that a column was skipped.

Reproduction

With polyglot-sql = 0.6.1 and current main (d5fd24eec5efaa4444eed6e5f009044214ccdc84):

use std::collections::HashMap;

use polyglot_sql::traversal::ExpressionWalk;
use polyglot_sql::{parse_by_name, rename_columns, Expression};

fn main() {
    let mapping = HashMap::from([("c".to_string(), "renamed".to_string())]);
    for (dialect, sql) in [
        ("oracle", "SELECT UPPER(t.c) FROM t"),
        ("oracle", "SELECT INSTR(t.c, 'x') FROM t"),
        ("oracle", "SELECT NVL2(t.c, 1, 0) FROM t"),
        ("oracle", "SELECT LAST_DAY(t.c) FROM t"),
        ("snowflake", "SELECT YEAR(t.c) FROM t"),
        ("bigquery", "SELECT STARTS_WITH(t.c, 'a') FROM t"),
    ] {
        let expr = parse_by_name(sql, dialect).unwrap().remove(0);
        let seen = expr.find_all(|e| matches!(e, Expression::Column(_))).len();
        let out = rename_columns(expr, &mapping);
        let sql_out = polyglot_sql::generate_by_name(&out, dialect).unwrap();
        println!("find_all columns={seen}  rename_columns -> {sql_out}");
    }
}

Output:

find_all columns=1  rename_columns -> SELECT UPPER(t.renamed) FROM t
find_all columns=1  rename_columns -> SELECT INSTR(t.c, 'x') FROM t
find_all columns=1  rename_columns -> SELECT NVL2(t.c, 1, 0) FROM t
find_all columns=1  rename_columns -> SELECT LAST_DAY(t.c) FROM t
find_all columns=1  rename_columns -> SELECT YEAR(t.c) FROM t
find_all columns=1  rename_columns -> SELECT STARTS_WITH(t.c, 'a') FROM t

find_all counts the column in every statement, so the parser sees it. The mutable walk is what skips it. The same happens with NEXT_DAY, ADD_MONTHS, MONTHS_BETWEEN, TO_NUMBER, INITCAP, FIRST_VALUE, JSON_VALUE, ARRAY_LENGTH, CONTAINS, ENDS_WITH, Postgres ~, IS DISTINCT FROM and ClickHouse's quantile(0.5)(x).

Comparison with sqlglot

sqlglot (30.11.0) renames all six. Its Expression.transform walks each node's args and does not care what type the node is:

import sqlglot
from sqlglot import exp

def rename(node):
    if isinstance(node, exp.Column) and node.name == "c":
        node.set("this", exp.to_identifier("renamed"))
    return node

for dialect, sql in [
    ("oracle", "SELECT UPPER(t.c) FROM t"),
    ("oracle", "SELECT INSTR(t.c, 'x') FROM t"),
    ("oracle", "SELECT NVL2(t.c, 1, 0) FROM t"),
    ("oracle", "SELECT LAST_DAY(t.c) FROM t"),
    ("snowflake", "SELECT YEAR(t.c) FROM t"),
    ("bigquery", "SELECT STARTS_WITH(t.c, 'a') FROM t"),
]:
    tree = sqlglot.parse_one(sql, read=dialect)
    print(tree.transform(rename).sql(dialect=dialect))

Output:

SELECT UPPER(t.renamed) FROM t
SELECT INSTR(t.renamed, 'x') FROM t
SELECT NVL2(t.renamed, 1, 0) FROM t
SELECT LAST_DAY(t.renamed) FROM t
SELECT YEAR(t.renamed) FROM t
SELECT STARTS_WITH(t.renamed, 'a') FROM t

Problem

transform_map and transform both end up in transform_recursive_inner. That function only uses the derived child visitor (for_each_child_mut) for the node kinds listed in uses_generated_dispatch. Every other node goes to transform_recursive_reference, and its catch-all arm (other => other) hands the node back without visiting its children.

The read-only walkers (find_all, dfs, contains) use the derived visitor for every node, which is why they see what the mutable walk misses.

Everything in ast_transforms (rename_columns, rename_tables, qualify_columns, replace_nodes, ...) is built on transform, so it all inherits this. #358 was the same bug for ALL / ANY subqueries.

Requirement

A mutable traversal that visits every child of every node kind. Either of these would work:

• A traversal::transform_all(expr, f): the explicit-stack loop from transform_recursive_inner, but using for_each_child_mut for every node. It stays iterative, and transform_map keeps its current behaviour for the dialect transforms that depend on it.
• Or make ast_children::for_each_child_mut and for_each_child public, so callers can write their own walk.

The ast_transforms helpers should then use the complete walk, since they are meant to reach every node.

References

• polyglot/crates/polyglot-sql/src/dialects/mod.rs

  Line 689 in d5fd24e

  fn uses_generated_dispatch(expression: &Expression) -> bool {

  (uses_generated_dispatch)

• polyglot/crates/polyglot-sql/src/dialects/mod.rs

  Line 2291 in d5fd24e

  other => other,

  (catch-all arm of transform_recursive_reference)

• polyglot/crates/polyglot-sql/src/ast_children.rs

  Line 44 in d5fd24e

  pub(crate) fn for_each_child_mut(

  (for_each_child_mut, crate-private)
• traversal::transform_map does not recurse into ALL/ANY quantified subqueries #358

[tobilg]

Check whether this still persists in version 0.13.0 please, thanks!