Update SqlTableAliasSimplifier to handle CTEs (#1516)

Similar to the updates for the other optimizers in
#1503 and
#1504, this PR implements the
CTE case for the `SqlTableAliasSimplifier`

metricflow/sql/optimizer/table_alias_simplifier.py

@@ -26,11 +26,11 @@ class SqlTableAliasSimplifierVisitor(SqlQueryPlanNodeVisitor[SqlQueryPlanNode]):
 
     @override
     def visit_cte_node(self, node: SqlCteNode) -> SqlQueryPlanNode:
-        raise NotImplementedError
+        return node.with_new_select(node.select_statement.accept(self))
 
     def visit_select_statement_node(self, node: SqlSelectStatementNode) -> SqlQueryPlanNode:  # noqa: D102
-        # If there is only a single parent, no table aliases are required since there's no ambiguity.
-        should_simplify_table_aliases = len(node.parent_nodes) <= 1
+        # If there is only a single source in the SELECT, no table aliases are required since there's no ambiguity.
+        should_simplify_table_aliases = len(node.join_descs) == 0
 
         if should_simplify_table_aliases:
             return SqlSelectStatementNode.create(
@@ -41,6 +41,10 @@ def visit_select_statement_node(self, node: SqlSelectStatementNode) -> SqlQueryP
                 ),
                 from_source=node.from_source.accept(self),
                 from_source_alias=node.from_source_alias,
+                cte_sources=tuple(
+                    cte_source.with_new_select(cte_source.select_statement.accept(self))
+                    for cte_source in node.cte_sources
+                ),
                 group_bys=tuple(
                     SqlSelectColumn(expr=x.expr.rewrite(should_render_table_alias=False), column_alias=x.column_alias)
                     for x in node.group_bys
@@ -59,6 +63,9 @@ def visit_select_statement_node(self, node: SqlSelectStatementNode) -> SqlQueryP
             select_columns=node.select_columns,
             from_source=node.from_source.accept(self),
             from_source_alias=node.from_source_alias,
+            cte_sources=tuple(
+                cte_source.with_new_select(cte_source.select_statement.accept(self)) for cte_source in node.cte_sources
+            ),
             join_descs=tuple(
                 SqlJoinDescription(
                     right_source=x.right_source.accept(self),

tests_metricflow/snapshots/test_cte_table_alias_simplifier.py/str/test_table_alias_no_simplification__result.txt

@@ -0,0 +1,41 @@
+test_name: test_table_alias_no_simplification
+test_filename: test_cte_table_alias_simplifier.py
+docstring:
+  Tests that table aliases in the SELECT statement of a CTE are not removed when required.
+---
+optimizer:
+  SqlTableAliasSimplifier
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      from_source_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table_0 from_source_alias
+    INNER JOIN
+      test_schema.test_table_1 right_source_alias
+    ON
+      from_source_alias.col_0 = right_source_alias.col_0
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      from_source_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table_0 from_source_alias
+    INNER JOIN
+      test_schema.test_table_1 right_source_alias
+    ON
+      from_source_alias.col_0 = right_source_alias.col_0
+  )
+
+  SELECT
+    cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias

tests_metricflow/snapshots/test_cte_table_alias_simplifier.py/str/test_table_alias_simplification__result.txt

@@ -0,0 +1,85 @@
+test_name: test_table_alias_simplification
+test_filename: test_cte_table_alias_simplifier.py
+docstring:
+  Tests that table aliases in the SELECT statement of a CTE are removed when not needed.
+---
+optimizer:
+  SqlTableAliasSimplifier
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+      , test_table_alias.col_1 AS cte_source_0__col_1
+    FROM (
+      -- CTE source 0 sub-query
+      SELECT
+        test_table_alias.col_0 AS cte_source_0_subquery__col_0
+        , test_table_alias.col_0 AS cte_source_0_subquery__col_1
+      FROM test_schema.test_table test_table_alias
+    ) cte_source_0_subquery
+  )
+
+  , cte_source_1 AS (
+    -- CTE source 1
+    SELECT
+      test_table_alias.col_0 AS cte_source_1__col_0
+      , test_table_alias.col_1 AS cte_source_1__col_1
+    FROM (
+      -- CTE source 1 sub-query
+      SELECT
+        cte_source_0_alias.cte_source_0__col_0 AS cte_source_1_subquery__col_0
+        , cte_source_0_alias.cte_source_0__col_0 AS cte_source_1_subquery__col_1
+      FROM cte_source_0 cte_source_0_alias
+    ) cte_source_1_subquery
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+    , right_source_alias.right_source__col_1 AS top_level__col_1
+  FROM cte_source_0 cte_source_0_alias
+  INNER JOIN
+    cte_source_1 right_source_alias
+  ON
+    cte_source_0_alias.cte_source_0__col_1 = right_source_alias.right_source__col_1
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      col_0 AS cte_source_0__col_0
+      , col_1 AS cte_source_0__col_1
+    FROM (
+      -- CTE source 0 sub-query
+      SELECT
+        col_0 AS cte_source_0_subquery__col_0
+        , col_0 AS cte_source_0_subquery__col_1
+      FROM test_schema.test_table test_table_alias
+    ) cte_source_0_subquery
+  )
+
+  , cte_source_1 AS (
+    -- CTE source 1
+    SELECT
+      col_0 AS cte_source_1__col_0
+      , col_1 AS cte_source_1__col_1
+    FROM (
+      -- CTE source 1 sub-query
+      SELECT
+        cte_source_0__col_0 AS cte_source_1_subquery__col_0
+        , cte_source_0__col_0 AS cte_source_1_subquery__col_1
+      FROM cte_source_0 cte_source_0_alias
+    ) cte_source_1_subquery
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+    , right_source_alias.right_source__col_1 AS top_level__col_1
+  FROM cte_source_0 cte_source_0_alias
+  INNER JOIN
+    cte_source_1 right_source_alias
+  ON
+    cte_source_0_alias.cte_source_0__col_1 = right_source_alias.right_source__col_1

tests_metricflow/snapshots/test_table_alias_simplifier.py/SqlQueryPlan/test_table_alias_simplification__after_alias_simplification.sql

@@ -1,7 +1,7 @@
 test_name: test_table_alias_simplification
 test_filename: test_table_alias_simplifier.py
 docstring:
-  Tests a case where no pruning should occur.
+  Tests that table aliases are removed when not needed.
 ---
 -- test0
 SELECT

tests_metricflow/snapshots/test_table_alias_simplifier.py/SqlQueryPlan/test_table_alias_simplification__before_alias_simplification.sql

@@ -1,7 +1,7 @@
 test_name: test_table_alias_simplification
 test_filename: test_table_alias_simplifier.py
 docstring:
-  Tests a case where no pruning should occur.
+  Tests that table aliases are removed when not needed.
 ---
 -- test0
 SELECT