From 8bdc496581a942b95f720eb586874050b991254e Mon Sep 17 00:00:00 2001 From: Maciej Brencz Date: Wed, 11 Sep 2024 11:59:24 +0200 Subject: [PATCH] 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 --- sql_metadata/parser.py | 6 +++++- test/test_getting_columns.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index d003c547..a8cac3e6 100644 --- a/sql_metadata/parser.py +++ b/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] diff --git a/test/test_getting_columns.py b/test/test_getting_columns.py index 8f6c20b5..74531760 100644 --- a/test/test_getting_columns.py +++ b/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():