diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index ed35226d..3fd82f2d 100644 --- a/sql_metadata/parser.py +++ b/sql_metadata/parser.py @@ -550,19 +550,11 @@ def with_names(self) -> List[str]: with_names.append(prev_token.left_expanded) else: with_names.append(token.left_expanded) - # move to next with if exists, this with ends with - # ) + , if many withs or ) + SELECT if one - # need to move to next as AS can be in - # sub-queries inside with definition - while token.next_token and not ( - token.is_right_parenthesis - and ( - token.next_token.is_punctuation - or token.next_token.normalized in WITH_ENDING_KEYWORDS - ) - ): + # move to next with query end + while token.next_token and not token.is_with_query_end: token = token.next_token if token.next_token.normalized in WITH_ENDING_KEYWORDS: + # end of with block self._is_in_with_block = False else: token = token.next_token diff --git a/test/test_with_statements.py b/test/test_with_statements.py index 64371196..6ad04c9f 100644 --- a/test/test_with_statements.py +++ b/test/test_with_statements.py @@ -90,3 +90,55 @@ def test_multiple_with_statements_with_with_columns(): assert parser.columns_aliases_names == ["c1", "c2", "c3", "c4"] assert parser.columns_aliases == {"c1": "*", "c2": "*", "c3": "c5", "c4": "c6"} assert parser.query_type == QueryType.SELECT + + +def test_complicated_with(): + query = """ + WITH uisd_filter_table as ( + select + session_id, + srch_id, + srch_ci, + srch_co, + srch_los, + srch_sort_type, + impr_list + from + uisd + where + datem <= date_sub(date_add(current_date(), 92), 7 * 52) + and lower(srch_sort_type) in ('expertpicks', 'recommended') + and srch_ci <= date_sub(date_add(current_date(), 92), 7 * 52) + and srch_co >= date_sub(date_add(current_date(), 1), 7 * 52) + ) + select + DISTINCT session_id, + srch_id, + srch_ci, + srch_co, + srch_los, + srch_sort_type, + l.impr_property_id as expe_property_id, + l.impr_position_across_pages + from + uisd_filter_table lateral view explode(impr_list) table as l + """ + parser = Parser(query) + assert parser.query_type == QueryType.SELECT + assert parser.with_names == ["uisd_filter_table"] + assert parser.tables == [ + "uisd", + "impr_list", + ] # this one is wrong too should be table + assert parser.columns == [ + "session_id", + "srch_id", + "srch_ci", + "srch_co", + "srch_los", + "srch_sort_type", + "impr_list", + "datem", + "l.impr_property_id", + "l.impr_position_across_pages", + ]