Handle when column name wildcard is prefixed by table name #296 (#330)

* Handle when column name wildcard is prefixed by table name

* Add test for getting wildcard column with table prefix

* Formatting tests

* linter: _resolve_nested_query is complex, ignore for now

Co-authored-by: Dimas <dimasjkl@gmail.com>

Author: macbre

sql_metadata/parser.py

@@ -809,7 +809,8 @@ def _resolve_sub_queries(self, column: str) -> List[str]:
         return column if isinstance(column, list) else [column]
 
     @staticmethod
-    def _resolve_nested_query(
+    # pylint:disable=too-many-return-statements
+    def _resolve_nested_query(  # noqa: C901
         subquery_alias: str,
         nested_queries_names: List[str],
         nested_queries: Dict,
@@ -845,6 +846,9 @@ def _resolve_nested_query(
             # handle case when column name is used but subquery select all by wildcard
             if "*" in subparser.columns:
                 return column_name
+            for table in subparser.tables:
+                if f"{table}.*" in subparser.columns:
+                    return column_name
             raise exc  # pragma: no cover
         resolved_column = subparser.columns[column_index]
         return [resolved_column]

test/test_getting_columns.py

@@ -92,6 +92,15 @@ def test_getting_columns():
         "test",
     ]
     assert Parser("SELECT /* a comment */ bar FROM test_table").columns == ["bar"]
+    assert (
+        Parser(
+            """
+                WITH foo AS (SELECT test_table.* FROM test_table)
+                SELECT foo.bar FROM foo
+            """
+        ).columns
+        == ["test_table.*", "bar"]
+    )
 
 
 def test_columns_with_order_by():