rule: AM07

test_pass_known_number_of_result_columns_1:
  pass_str: |
    select a, b
    from t
    union all
    select c, d
    from k

test_fail_known_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_known_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

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_known_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 t2
    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_known_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_fail_cte_no_select_final_statement:
  fail_str:
    UPDATE sometable
    SET
    sometable.baz = mycte.bar
    FROM
    (SELECT
    foo,
    bar
    FROM mytable1
    
    UNION ALL
    
    SELECT
    bar
    FROM mytable2) as k

test_tsql_select_system_as_identifier:
  pass_str: |
    SELECT @@IDENTITY AS 'Identity'
  configs:
    core:
      dialect: tsql