fix: CTE queries with non-SELECT statements (#25014)

This commit is contained in:
Daniel Vaz Gaspar
2023-08-19 15:49:15 +01:00
committed by GitHub
parent 6b660c86a4
commit 357986103b
2 changed files with 136 additions and 0 deletions

View File

@@ -1029,6 +1029,87 @@ FROM foo f"""
assert sql.is_select()
def test_cte_insert_is_not_select() -> None:
"""
Some CTEs with lowercase select are not correctly identified as SELECTS.
"""
sql = ParsedQuery(
"""WITH foo AS(
INSERT INTO foo (id) VALUES (1) RETURNING 1
) select * FROM foo f"""
)
assert sql.is_select() is False
def test_cte_delete_is_not_select() -> None:
"""
Some CTEs with lowercase select are not correctly identified as SELECTS.
"""
sql = ParsedQuery(
"""WITH foo AS(
DELETE FROM foo RETURNING *
) select * FROM foo f"""
)
assert sql.is_select() is False
def test_cte_is_not_select_lowercase() -> None:
"""
Some CTEs with lowercase select are not correctly identified as SELECTS.
"""
sql = ParsedQuery(
"""WITH foo AS(
insert into foo (id) values (1) RETURNING 1
) select * FROM foo f"""
)
assert sql.is_select() is False
def test_cte_with_multiple_selects() -> None:
sql = ParsedQuery(
"WITH a AS ( select * from foo1 ), b as (select * from foo2) SELECT * FROM a;"
)
assert sql.is_select()
def test_cte_with_multiple_with_non_select() -> None:
sql = ParsedQuery(
"""WITH a AS (
select * from foo1
), b as (
update foo2 set id=2
) SELECT * FROM a"""
)
assert sql.is_select() is False
sql = ParsedQuery(
"""WITH a AS (
update foo2 set name=2
),
b as (
select * from foo1
) SELECT * FROM a"""
)
assert sql.is_select() is False
sql = ParsedQuery(
"""WITH a AS (
update foo2 set name=2
),
b as (
update foo1 set name=2
) SELECT * FROM a"""
)
assert sql.is_select() is False
sql = ParsedQuery(
"""WITH a AS (
INSERT INTO foo (id) VALUES (1)
),
b as (
select 1
) SELECT * FROM a"""
)
assert sql.is_select() is False
def test_unknown_select() -> None:
"""
Test that `is_select` works when sqlparse fails to identify the type.