fix: CTE queries with non-SELECT statements

[dpgaspar]

SUMMARY

Improves SQLLab parse check for non SELECT statements for CTE.

The following statement was being accepted as a valid SELECT statement:

WITH a AS ( INSERT INTO foo (id) VALUES (1) RETURNING id ) SELECT * FROM a;

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

• Has associated issue:
• Required feature flags:
• Changes UI
• Includes DB Migration (follow approval process in SIP-59)
  • Migration is atomic, supports rollback & is backwards-compatible
  • Confirm DB migration upgrade and downgrade tested
  • Runtime estimates and downtime expectations provided
• Introduces new feature or API
• Removes existing feature or API

[betodealmeida]

Looks good, but I'm worried about all the possible edge case we're missing.

I wonder if we should also use sqloxide to prevent false negatives. We'd try to parse with sqloxide, check for non-selects, return false if we do. If we don't, we then run the current flow with sqlparse.

>>> parse_sql("WITH a AS ( INSERT INTO foo (id) VALUES (1) RETURNING id ) SELECT * FROM a;", "ansi")
[{'Query': {'with': {'recursive': False, 'cte_tables': [{'alias': {'name': {'value': 'a', 'quote_style': None}, 'columns': []}, 'query': {'with': None, 'body': {'Insert': {'Insert': {'or': None, 'into': True, 'table_name': [{'value': 'foo', 'quote_style': None}], 'columns': [{'value': 'id', 'quote_style': None}], 'overwrite': False, 'source': {'with': None, 'body': {'Values': {'explicit_row': False, 'rows': [[{'Value': {'Number': ('1', False)}}]]}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}, 'partitioned': None, 'after_columns': [], 'table': False, 'on': None, 'returning': [{'UnnamedExpr': {'Identifier': {'value': 'id', 'quote_style': None}}}]}}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}, 'from': None}]}, 'body': {'Select': {'distinct': False, 'top': None, 'projection': [{'Wildcard': {'opt_exclude': None, 'opt_except': None, 'opt_rename': None, 'opt_replace': None}}], 'into': None, 'from': [{'relation': {'Table': {'name': [{'value': 'a', 'quote_style': None}], 'alias': None, 'args': None, 'with_hints': []}}, 'joins': []}], 'lateral_views': [], 'selection': None, 'group_by': [], 'cluster_by': [], 'distribute_by': [], 'sort_by': [], 'having': None, 'qualify': None}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}}]

[dpgaspar]

Looks good, but I'm worried about all the possible edge case we're missing.

I wonder if we should also use sqloxide to prevent false negatives. We'd try to parse with sqloxide, check for non-selects, return false if we do. If we don't, we then run the current flow with sqlparse.

>>> parse_sql("WITH a AS ( INSERT INTO foo (id) VALUES (1) RETURNING id ) SELECT * FROM a;", "ansi")
[{'Query': {'with': {'recursive': False, 'cte_tables': [{'alias': {'name': {'value': 'a', 'quote_style': None}, 'columns': []}, 'query': {'with': None, 'body': {'Insert': {'Insert': {'or': None, 'into': True, 'table_name': [{'value': 'foo', 'quote_style': None}], 'columns': [{'value': 'id', 'quote_style': None}], 'overwrite': False, 'source': {'with': None, 'body': {'Values': {'explicit_row': False, 'rows': [[{'Value': {'Number': ('1', False)}}]]}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}, 'partitioned': None, 'after_columns': [], 'table': False, 'on': None, 'returning': [{'UnnamedExpr': {'Identifier': {'value': 'id', 'quote_style': None}}}]}}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}, 'from': None}]}, 'body': {'Select': {'distinct': False, 'top': None, 'projection': [{'Wildcard': {'opt_exclude': None, 'opt_except': None, 'opt_rename': None, 'opt_replace': None}}], 'into': None, 'from': [{'relation': {'Table': {'name': [{'value': 'a', 'quote_style': None}], 'alias': None, 'args': None, 'with_hints': []}}, 'joins': []}], 'lateral_views': [], 'selection': None, 'group_by': [], 'cluster_by': [], 'distribute_by': [], 'sort_by': [], 'having': None, 'qualify': None}}, 'order_by': [], 'limit': None, 'offset': None, 'fetch': None, 'locks': []}}]

nice!, will try that!

[betodealmeida]

<3