Repository: sqlfluff/sqlfluff
Pull request: Rule: check number of columns in set operations match
<!--Thanks for adding this feature!-->

<!--Please give the Pull Request a meaningful title for the release notes-->

### Brief summary of the change made
<!--Please include `fixes #XXXX` to automatically close any corresponding issue when the pull request is merged. Alternatively if not fully closed you can say `makes progress on #XXXX`.-->
Fixes https://github.com/sqlfluff/sqlfluff/issues/1741

### Are there any other side effects of this change that we should be aware of?

Somewhat related to L044 (check that wildcards resolve to columns), however, if wildcard columns can't be resolved, doesn't fail, leaves that check to L044


### Pull Request checklist
- [ ] Please confirm you have completed any of the necessary steps below.

- Included test cases to demonstrate any code changes, which may be one or more of the following:
  - `.yml` rule test cases in `test/fixtures/rules/std_rule_cases`.
  - `.sql`/`.yml` parser test cases in `test/fixtures/dialects` (note YML files can be auto generated with `tox -e generate-fixture-yml`).
  - Full autofix test cases in `test/fixtures/linter/autofix`.
  - Other.
- Added appropriate documentation for the change.
- Created GitHub issues for any relevant followup/future enhancements if appropriate.


Problem statement:
Add rule to check number of columns in set operations
Add rule to check that both queries in a set operation contain the same number of columns.

For example:

```
SELECT a, b, c
FROM test_table

UNION ALL

SELECT a, b
FROM test_table_2
```

The above produces a compilation error in Snowflake (`"invalid number of result columns for set operator input branches, expected 3, got 2 in branch 1"`), however it passes linting and parsing. 

Seems similar to rule L044, and could raise a parsing error.

Hints from issue discussion:
I like the idea. I suggest that it should raise a linting error, not a parsing error. For technical reasons, rules don't raise parse errors. Parse errors are generally raised only by the parser, when the SQL does not obey the grammar (i.e. the dialect).
Good idea!
A clarification would be if one of the statements contains a wildcard we can't determine the number of columns, so let's not raise the lint error in that case as it would already be covered by L044.

Commit message: edit previous commit message

Patch to review:
diff --git a/src/sqlfluff/rules/L070.py b/src/sqlfluff/rules/L070.py
new file mode 100644
index 00000000000..34c300536e3
--- /dev/null
+++ b/src/sqlfluff/rules/L070.py
@@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
+
+    def _find_all_ctes(self, context: RuleContext, cte_dict):
+        for i in range(len(context.parent_stack)):
+            if context.parent_stack[i].type != "common_table_expression":
+                parent_query = SelectCrawler(
+                    context.parent_stack[i], context.dialect
+                ).query_tree
+                cte_dict.update(self._find_all_ctes_utils(parent_query, cte_dict))
+
+        return cte_dict
+
+    def __resolve_wildcard(
+        self,
+        context: RuleContext,
+        query,
+        parent_query,
+        resolve_targets,
+        all_ctes,
+    ):
+        """Attempt to resolve the wildcard to a list of selectables."""
+        process_queries = query.selectables
+        # if one of the source queries for a query within the set is a
+        # set expression, just use the first query. If that first query isn't
+        # reflective of the others, that will be caught when that segment
+        # is processed
+        if query.selectables[0].parent:
+            if query.selectables[0].parent.type == "set_expression":
+                process_queries = [query.selectables[0]]
+
+        for selectable in process_queries:
+            if selectable.get_wildcard_info():
+                for wildcard in selectable.get_wildcard_info():
+                    if wildcard.tables:
+                        for wildcard_table in wildcard.tables:
+                            alias_info = selectable.find_alias(wildcard_table)
+                            # attempt to resolve alias or table name to a cte;
+                            if alias_info:
+                                select_info_target = SelectCrawler.get(
+                                    parent_query, alias_info.from_expression_element
+                                )[0]
+                                if isinstance(select_info_target, str):
+                                    cte = None
+                                    # handles case where cte is in subquery
+                                    cte_child = query.lookup_cte(select_info_target)
+                                    if select_info_target.upper() in all_ctes:
+                                        cte = all_ctes[select_info_target.upper()]
+                                    # check parent cte
+                                    if cte:
+                                        self.__resolve_wildcard(
+                                            context,
+                                            cte,
+                                            parent_query,
+                                            resolve_targets,
+                                            all_ctes,
+                                        )
+                                    elif cte_child:
+                                        self.__resolve_wildcard(
+                                            context,
+                                            cte_child,
+                                            parent_query,
+                                            resolve_targets,
+                                            all_ctes,
+                                        )
+                                    else:
+                                        resolve_targets.append(wildcard)
+
+                                # if select_info_target is not a string
+                                # , process as a subquery
+                                else:
+                                    self.__resolve_wildcard(
+                                        context,
+                                        select_info_target,
+                                        parent_query,
+                                        resolve_targets,
+                                        all_ctes,
+                                    )
+                            else:
+                                cte = None
+                                cte_child = query.lookup_cte(wildcard_table)
+                                if wildcard_table.upper() in all_ctes:
+                                    cte = all_ctes[wildcard_table.upper()]
+                                if cte:
+                                    self.__resolve_wildcard(
+                                        context,
+                                        cte,
+                                        parent_query,
+                                        resolve_targets,
+                                        all_ctes,
+                                    )
+                                elif cte_child:
+                                    self.__resolve_wildcard(
+                                        context,
+                                        cte_child,
+                                        parent_query,
+                                        resolve_targets,
+                                        all_ctes,
+                                    )
+                                else:
+                                    resolve_targets.append(wildcard)
+                    # if there is no table specified, it is likely a subquery
+                    else:
+                        query_list = SelectCrawler.get(
+                            query, query.selectables[0].selectable
+                        )
+                        for o in query_list:
+                            if isinstance(o, Query):
+                                self.__resolve_wildcard(
+                                    context, o, parent_query, resolve_targets, all_ctes
+                                )
+                                return resolve_targets
+            else:
+                resolve_targets.extend(
+                    [target for target in selectable.select_info.select_targets]
+                )
+
+        return resolve_targets
+
+    def _get_select_target_counts(self, context: RuleContext, crawler):
+        """Given a set expression, get the number of select targets in each query."""
+        select_list = None
+        select_target_counts = set()
+        set_selectables = crawler.query_tree.selectables
+        resolved_wildcard = True
+        parent_crawler = SelectCrawler(context.parent_stack[0], context.dialect)
+        # for each selectable in the set
+        # , add the number of select targets to a set; in the end
+        # length of the set should be one (one size)
+        for selectable in set_selectables:
+            # if the set query contains a wildcard
+            # , attempt to resolve wildcard to a list of
+            # select targets that can be counted
+            if selectable.get_wildcard_info():
+                # to start, get a list of all of the ctes in the parent
+                # stack to check whether they resolve to wildcards
+
+                all_cte_queries = self._find_all_ctes(context, {})
+                select_crawler = SelectCrawler(
+                    selectable.selectable, context.dialect
+                ).query_tree
+
+                select_list = self.__resolve_wildcard(
+                    context,
+                    select_crawler,
+                    parent_crawler.query_tree,
+                    [],
+                    all_cte_queries,
+                )
+
+                # get the number of resolved targets plus the total number of
+                # targets minus the number of wildcards
+                # if all wildcards have been resolved this adds up to the
+                # total number of select targets in the query
+                select_target_counts.add(
+                    len(select_list)
+                    + (
+                        len(selectable.select_info.select_targets)
+                        - len(selectable.get_wildcard_info())
+                    )
+                )
+                for select in select_list:
+                    if isinstance(select, WildcardInfo):
+                        resolved_wildcard = False
+            else:
+                # if there is no wildcard in the query use the count of select targets
+                select_list = selectable.select_info.select_targets
+                select_target_counts.add(len(select_list))
+
+        return (select_target_counts, resolved_wildcard)
+
+    def _eval(self, context: RuleContext) -> Optional[LintResult]:
+        """All queries in set expression should return the same number of columns."""
+        assert context.segment.is_type("set_expression")
+        crawler = SelectCrawler(context.segment, context.dialect)
+
+        set_segment_select_sizes, resolve_wildcard = self._get_select_target_counts(
+            context, crawler
+        )
+        # if queries had different select target counts
+        # and all wildcards have been resolved; fail
+        if len(set_segment_select_sizes) > 1 and resolve_wildcard:
+            return LintResult(anchor=context.segment)
+
+        return LintResult()
diff --git a/test/fixtures/rules/std_rule_cases/L070.yml b/test/fixtures/rules/std_rule_cases/L070.yml
new file mode 100644
index 00000000000..2577e4bef4b
--- /dev/null
+++ b/test/fixtures/rules/std_rule_cases/L070.yml
@@ -0,0 +1,482 @@
+rule: L070
+
+test_pass_known_number_of_result_columns_1:
+  pass_str: |
+    select a, b
+    from t
+    union all
+    select c, d
+    from k
+
+test_fail_unknown_number_of_result_columns_1:
+  fail_str: |
+    select a
+    from t
+    union all
+    select c, d
+    from k
+
+test_pass_known_number_of_result_columns_2:
+  desc: Columns are specified in CTE so * in final query will return only columns specified earlier.
+  pass_str: |
+    with cte as (
+        select
+            a, b
+        from
+            t
+    )
+    select * from cte
+    union
+    select c, d
+    from t2
+
+test_fail_unknown_number_of_result_columns_2:
+  fail_str: |
+    with cte as (
+        select
+            a, b, c
+        from
+            t
+    )
+    select *
+    from cte
+    union
+    select d, e
+    from t
+
+# pass3;
+# fail3;
+test_pass_known_number_of_result_columns_3:
+  pass_str: |
+    with cte as (
+        select
+            *
+        from
+            t
+    )
+    select a, b from cte
+
+test_pass_known_number_of_result_columns_4:
+  pass_str: |
+    with cte1 as (
+        select
+            *
+        from
+            t
+    ),
+    cte2 as (
+        select
+            a, b
+        from
+            t
+    )
+    select
+        cte1.a
+        , cte1.d
+        , cte2.*
+    from cte1
+    join cte2
+    using (a)
+
+    union
+    select e, f, g, h
+    from cte3
+
+test_fail_unknown_number_of_result_columns_3:
+  fail_str: |
+    with cte1 as (
+        select
+            *
+        from
+            t
+    ),
+    cte2 as (
+        select
+            a, b
+        from
+            t
+    )
+    select
+        cte1.a,
+        cte2.*
+    from cte1
+    join cte2
+    using (a)
+    union
+    select e, f
+    from cte3
+
+test_pass_known_number_of_result_columns_5:
+  pass_str: |
+    select * from (
+        select
+            a, b
+        from
+            t
+    )
+    union
+    select c, d
+    from t2
+
+
+test_pass_known_number_of_result_columns_6:
+  pass_str: |
+    with cte2 as (
+      select b
+      from t1
+    )
+    select
+        t_alias.*
+    from cte1
+    join (select a from t) as t_alias
+    using (a)
+    union
+    select *
+    from cte2
+
+test_fail_unknown_number_of_result_columns_4:
+  fail_str: |
+    select
+        t_alias.*
+    from t1
+    join (select a from t) as t_alias
+    using (a)
+    union
+    select a,b
+    from t2
+
+
+test_pass_known_number_of_result_columns_7:
+  pass_str: |
+    select a from t1
+    union all
+    select b from t2
+
+test_pass_unknown_wildcard_number_of_result_columns_8:
+  pass_str: |
+    select a from t1
+    union all
+    select * from t2
+
+test_pass_known_number_of_result_columns_9:
+  pass_str: |
+    with cte as (
+      select a from t1
+      union all
+      select b from t2
+    )
+
+    select * from cte
+    union all
+    select b from t3
+
+test_pass_known_number_of_result_columns_10:
+  desc: Columns are specified in cte_orders's body so cte_orders.* will return only these columns.
+  pass_str: |
+    WITH cte_orders AS (
+      SELECT customer_id, total
+      FROM orders
+    )
+
+    SELECT customers.name, cte_orders.*
+    FROM
+      customers,
+      cte_orders
+    WHERE
+      clients.id = orders.clientId
+
+test_pass_known_number_of_result_columns_11:
+  pass_str: select a, b from `d.t`
+            union all
+            select c, d from `d.t`
+  configs:
+    core:
+      dialect: bigquery
+
+test_fail_unknown_number_of_result_columns_5:
+  fail_str: select a, b, c from `d.t` union all select c, d from `d.t`
+  configs:
+    core:
+      dialect: bigquery
+
+test_pass_known_number_of_result_columns_13:
+  # Issue 915: Crash on statements that don't have a SELECT
+  pass_str: CREATE TABLE my_table (id INTEGER)
+
+test_pass_known_number_of_result_columns_14:
+  pass_str: |
+    SELECT *
+    FROM
+    (
+      WITH cte2 AS (
+        SELECT a, b
+        FROM table2
+      )
+      SELECT *
+      from cte2 as cte_al
+      UNION
+      SELECT e, f
+      FROM table3
+    )
+    UNION
+    SELECT k, l
+    FROM table4
+
+test_pass_known_number_of_result_columns_15:
+  pass_str: |
+    SELECT *
+    FROM
+    (
+      WITH cte2 AS (
+        SELECT *
+        FROM table2
+      )
+      SELECT *
+      from cte2 as cte_al
+      UNION
+      SELECT e, f
+      FROM table3
+    )
+    UNION
+    SELECT k, l
+    FROM table4
+
+test_fail_unknown_number_of_result_columns_6:
+  fail_str: |
+    SELECT *
+    FROM
+    (
+      WITH cte2 AS (
+        SELECT a
+        FROM table2
+      )
+      SELECT *
+      from cte2 as cte_al
+      UNION
+      SELECT e, f
+      FROM table3
+    )
+    UNION
+    SELECT k, l
+    FROM table4
+
+test_pass_known_number_of_result_columns_16:
+  pass_str: |
+    SELECT *
+    FROM
+    (
+      WITH cte2 AS (
+        SELECT a, b
+        FROM table2
+      )
+      SELECT *
+      from cte2 as cte_al
+    )
+    UNION
+    SELECT e, f
+    FROM table3
+
+test_pass_known_number_of_result_columns_17:
+  pass_str: |
+    SELECT *
+    FROM table1 as table2
+    UNION
+    SELECT e, f
+    FROM table3
+
+test_fail_unknown_number_of_result_columns_7:
+  fail_str: |
+    SELECT *
+    FROM
+    (
+      WITH cte2 AS (
+        SELECT a
+        FROM table2
+      )
+      SELECT *
+      from cte2 as cte_al
+    )
+    UNION
+    SELECT e, f
+    FROM table3
+
+test_pass_nested_ctes_1:
+  # Test for issue 1984
+  pass_str: |
+    with a as (
+        with b as (
+            select 1 from c
+        )
+        select * from b
+    )
+
+    select * from a
+    union all
+    select k
+    from t2
+
+test_fail_nested_ctes_1:
+  # Test for issue 1984
+  fail_str: |
+    with a as (
+        with b as (
+            select a from c
+        )
+
+        select * from b
+    )
+
+    select * from a
+    union
+    select a, b
+    from t2
+
+test_cte_reference_outer_2:
+  pass_str:
+    with a as (
+        select 1 from b
+    )
+
+    select * from (
+        select * from a
+    )
+
+    union
+
+    select 2
+    from c
+
+test_cte_tricky_nesting_3:
+  pass_str:
+    with b as (
+        select 1 from c
+    )
+
+    select * from (
+        with a as (
+            select * from b
+        )
+
+        select * from a
+    )
+
+    union
+    select a
+    from t2
+
+test_nested_and_same_level_ctes_4:
+  pass_str:
+    with a as (
+      with c as (
+        select 1 from d
+      ),
+
+      b as (
+        select * from c
+      )
+
+      select * from b
+    )
+
+    select * from a
+    union
+    select k
+    from t2
+
+test_nested_cte_references_outer_5:
+  pass_str:
+    with c as (
+      select 1 from d
+    ),
+
+    a as (
+      with b as (
+        select * from c
+      )
+
+      select * from b
+    )
+
+    select * from a
+    union
+    select k
+    from t2
+
+test_pass_join_inside_cte_with_unqualified:
+  pass_str:
+    with cte as (
+        select
+            *
+        from
+            t1
+        inner join t2
+    )
+
+    select a, b from cte
+    union
+    select c, d from cte2;
+
+test_pass_known_number_of_columns_in_two_join_subqueries:
+  pass_str:
+    select
+        *
+    from (
+        select a
+        from foo
+    ) t1
+    inner join (
+        select b
+        from bar
+    ) t2
+    union
+    select c, d from t3;
+
+test_fail_two_join_subqueries_one_with_unknown_number_of_columns:
+  fail_str:
+    select
+        *
+    from (
+        select b
+        from foo
+    ) t1
+    inner join (
+        select b
+        from bar
+    ) t2
+    union
+    select c, d, e from t3;
+
+
+test_query_on_snowflake_stage:
+  pass_str:
+    select mycolumn1
+    from @public.mytable1
+    union
+    select mycolumn2
+    from table2
+  configs:
+    core:
+      dialect: snowflake
+
+
+test_pass_cte_no_select_final_statement:
+  pass_str:
+    WITH mycte AS (
+      SELECT
+          foo,
+          bar
+      FROM mytable1
+    )
+
+    UPDATE sometable
+    SET
+        sometable.baz = mycte.bar
+    FROM
+        mycte;
+
+test_tsql_select_system_as_identifier:
+  pass_str: |
+    SELECT @@IDENTITY AS 'Identity'
+  configs:
+    core:
+      dialect: tsql


Code review comments:
diff_hunk: @@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
+
+    def _find_all_ctes(self, context: RuleContext, cte_dict):
+        for i in range(len(context.parent_stack)):
+            if context.parent_stack[i].type != "common_table_expression":
+                parent_query = SelectCrawler(
+                    context.parent_stack[i], context.dialect
+                ).query_tree
+                cte_dict.update(self._find_all_ctes_utils(parent_query, cte_dict))
+
+        return cte_dict
+
+    def __resolve_wildcard(
path: src/sqlfluff/rules/L070.py
text: @user1:
This function is really large and complex. Could you look at making it shorter by moving parts of it to separate helper functions?

diff_hunk: @@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
+
+    def _find_all_ctes(self, context: RuleContext, cte_dict):
+        for i in range(len(context.parent_stack)):
+            if context.parent_stack[i].type != "common_table_expression":
+                parent_query = SelectCrawler(
+                    context.parent_stack[i], context.dialect
path: src/sqlfluff/rules/L070.py
text: @user1:
```suggestion
    def _find_all_ctes(self, context: RuleContext, cte_dict):
        for seg in context.parent_stack:
                    if  not seg.is_type("common_table_expression"):
                        parent_query = SelectCrawler(
                            seg, context.dialect
```

diff_hunk: @@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
path: src/sqlfluff/rules/L070.py
text: @user1:
Nitpick: Since you're updating the dictionary in place, you don't need to return it from the function.

diff_hunk: @@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
+
+    def _find_all_ctes(self, context: RuleContext, cte_dict):
+        for i in range(len(context.parent_stack)):
+            if context.parent_stack[i].type != "common_table_expression":
+                parent_query = SelectCrawler(
+                    context.parent_stack[i], context.dialect
+                ).query_tree
+                cte_dict.update(self._find_all_ctes_utils(parent_query, cte_dict))
+
+        return cte_dict
+
+    def __resolve_wildcard(
+        self,
+        context: RuleContext,
+        query,
+        parent_query,
+        resolve_targets,
+        all_ctes,
+    ):
+        """Attempt to resolve the wildcard to a list of selectables."""
+        process_queries = query.selectables
+        # if one of the source queries for a query within the set is a
+        # set expression, just use the first query. If that first query isn't
+        # reflective of the others, that will be caught when that segment
+        # is processed
+        if query.selectables[0].parent:
+            if query.selectables[0].parent.type == "set_expression":
+                process_queries = [query.selectables[0]]
+
+        for selectable in process_queries:
+            if selectable.get_wildcard_info():
+                for wildcard in selectable.get_wildcard_info():
+                    if wildcard.tables:
+                        for wildcard_table in wildcard.tables:
+                            alias_info = selectable.find_alias(wildcard_table)
+                            # attempt to resolve alias or table name to a cte;
+                            if alias_info:
+                                select_info_target = SelectCrawler.get(
+                                    parent_query, alias_info.from_expression_element
+                                )[0]
+                                if isinstance(select_info_target, str):
+                                    cte = None
+                                    # handles case where cte is in subquery
+                                    cte_child = query.lookup_cte(select_info_target)
+                                    if select_info_target.upper() in all_ctes:
+                                        cte = all_ctes[select_info_target.upper()]
+                                    # check parent cte
+                                    if cte:
+                                        self.__resolve_wildcard(
+                                            context,
+                                            cte,
+                                            parent_query,
+                                            resolve_targets,
+                                            all_ctes,
+                                        )
+                                    elif cte_child:
+                                        self.__resolve_wildcard(
+                                            context,
+                                            cte_child,
+                                            parent_query,
+                                            resolve_targets,
+                                            all_ctes,
+                                        )
+                                    else:
+                                        resolve_targets.append(wildcard)
+
+                                # if select_info_target is not a string
+                                # , process as a subquery
+                                else:
+                                    self.__resolve_wildcard(
+                                        context,
+                                        select_info_target,
+                                        parent_query,
+                                        resolve_targets,
+                                        all_ctes,
+                                    )
+                            else:
+                                cte = None
+                                cte_child = query.lookup_cte(wildcard_table)
+                                if wildcard_table.upper() in all_ctes:
+                                    cte = all_ctes[wildcard_table.upper()]
path: src/sqlfluff/rules/L070.py
text: @user1:
```suggestion
                                cte = all_ctes.get(wildcard_table.upper())
                                cte_child = query.lookup_cte(wildcard_table)
```

diff_hunk: @@ -0,0 +1,249 @@
+"""Implementation of Rule L070."""
+from typing import Optional
+
+from sqlfluff.utils.analysis.select_crawler import Query, SelectCrawler, WildcardInfo
+from sqlfluff.core.rules import BaseRule, LintResult, RuleContext
+from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
+from sqlfluff.core.rules.doc_decorators import document_groups
+
+
+@document_groups
+class Rule_L070(BaseRule):
+    """Queries within set query produce different numbers of columns.
+
+    **Anti-pattern**
+
+    When writing set expressions
+    , all queries must return the
+    same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT
+                a,
+                b
+            FROM foo
+        )
+        SELECT * FROM cte
+        UNION
+        SELECT
+            c,
+            d,
+            e
+         FROM t
+
+    **Best practice**
+
+    Always specify columns when writing set queries
+    and ensure that they all seleect same number of columns
+
+    .. code-block:: sql
+
+        WITH cte AS (
+            SELECT a, b FROM foo
+        )
+        SELECT
+            a,
+            b
+        FROM cte
+        UNION
+        SELECT
+            c,
+            d
+        FROM t
+    """
+
+    groups = ("all",)
+    crawl_behaviour = SegmentSeekerCrawler({"set_expression"}, provide_raw_stack=True)
+
+    def _find_all_ctes_utils(self, query, cte_dict):
+        """Generate a list of all ctes in a query."""
+        # cte_dict = cte_dict | query.ctes
+        cte_dict.update(query.ctes)
+        for cte_name, cte in query.ctes.items():
+            cte_dict = self._find_all_ctes_utils(cte, cte_dict)
+        return cte_dict
+
+    def _find_all_ctes(self, context: RuleContext, cte_dict):
+        for i in range(len(context.parent_stack)):
+            if context.parent_stack[i].type != "common_table_expression":
+                parent_query = SelectCrawler(
+                    context.parent_stack[i], context.dialect
+                ).query_tree
+                cte_dict.update(self._find_all_ctes_utils(parent_query, cte_dict))
+
+        return cte_dict
+
+    def __resolve_wildcard(
+        self,
+        context: RuleContext,
+        query,
+        parent_query,
+        resolve_targets,
+        all_ctes,
+    ):
+        """Attempt to resolve the wildcard to a list of selectables."""
+        process_queries = query.selectables
+        # if one of the source queries for a query within the set is a
+        # set expression, just use the first query. If that first query isn't
+        # reflective of the others, that will be caught when that segment
+        # is processed
+        if query.selectables[0].parent:
+            if query.selectables[0].parent.type == "set_expression":
path: src/sqlfluff/rules/L070.py
text: @user1:
```suggestion
            if query.selectables[0].parent.is_type("set_expression"):
```