Merge pull request #162 from collerek/check_with

Fix parsing with clauses, fix for #161
macbre · Jun 10, 2021 · 17798bc · 17798bc
2 parents bbe8d25 + 7336e45
commit 17798bc
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 11 deletions.
diff --git 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
@@ -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",
+    ]