How to get CTE tables from AST?

[kasvith]

Hi,

Is there a way to collect all CTE tables from an ast?

WITH london1_monthly_revenue AS (
  SELECT
    EXTRACT(MONTH FROM date) as month,
    SUM(unit_price * quantity) AS revenue
  FROM sales
  WHERE EXTRACT(YEAR FROM date) = 2021
    AND branch = 'London-1'
  GROUP BY 1
),
london2_monthly_revenue AS (
  SELECT
    EXTRACT(MONTH FROM date) as month,
    SUM(unit_price * quantity) AS revenue
  FROM sales
  WHERE EXTRACT(YEAR FROM date) = 2021
    AND branch = 'London-2'
  GROUP BY 1
)
SELECT
  l1.month,
  l1.revenue + l2.revenue AS london_revenue,
  l1.revenue AS london1_revenue,
  l2.revenue AS london2_revenue
FROM london1_monthly_revenue l1, london2_monthly_revenue l2
WHERE l1.month = l2.month

I want to collect london1_monthly_revenue and london2_monthly_revenue from above expression

[kasvith]

I did something like this

    select.with.iter().for_each(|with| {
        with.cte_tables.iter().for_each(|cte_table| {
            cte_tables.insert(cte_table.alias.name.value.clone());
        });
    });

[alamb]

Maybe you want to look into Visitor: https://docs.rs/sqlparser/latest/sqlparser/ast/trait.Visitor.html#example

Which will handle the traversal for you

[kasvith]

hmm, what about the approach i used?

    select.with.iter().for_each(|with| {
        with.cte_tables.iter().for_each(|cte_table| {
            cte_tables.insert(cte_table.alias.name.value.clone());
        });
    });

[alamb]

If it works for you 👍

[kasvith]

it worked so far :)