Improved handling of custom SQL functions when getting the list of columns

sql_metadata/parser.py

@@ -57,7 +57,7 @@ def query(self) -> str:
     @property
     def tokens(self) -> List[SQLToken]:
         """
-        :rtype: list[SQLToken]
+        Tokenizes the query
         """
         if self._tokens is not None:
             return self._tokens
@@ -129,7 +129,7 @@ def tokens(self) -> List[SQLToken]:
     @property
     def columns(self) -> List[str]:
         """
-        :rtype: list[str]
+        Returns the list columns this query refers to
         """
         if self._columns is not None:
             return self._columns
@@ -144,10 +144,16 @@ def columns(self) -> List[str]:
                     token.last_keyword_normalized in KEYWORDS_BEFORE_COLUMNS
                     and token.previous_token.normalized != "AS"
                 ):
-                    if token.normalized not in FUNCTIONS_IGNORED and not (
-                        # aliases of sub-queries i.e.: select from (...) <alias>
-                        token.previous_token.is_right_parenthesis
-                        and token.value in subqueries_names
+                    if (
+                        token.normalized not in FUNCTIONS_IGNORED
+                        and not (
+                            # aliases of sub-queries i.e.: select from (...) <alias>
+                            token.previous_token.is_right_parenthesis
+                            and token.value in subqueries_names
+                        )
+                        # custom functions - they are followed by the parenthesis
+                        # e.g. custom_func(...
+                        and not token.next_token.is_left_parenthesis
                     ):
                         column = token.table_prefixed_column(tables_aliases)
                         self._add_to_columns_subsection(
@@ -204,7 +210,7 @@ def columns_dict(self) -> Dict[str, List[str]]:
     @property
     def tables(self) -> List[str]:
         """
-        :rtype: list[str]
+        Return the list of tables this query refers to
         """
         if self._tables is not None:
             return self._tables
@@ -247,8 +253,6 @@ def tables(self) -> List[str]:
     def limit_and_offset(self) -> Optional[Tuple[int, int]]:
         """
         Returns value for limit and offset if set
-
-        :rtype: (int, int)
         """
         if self._limit_and_offset is not None:
             return self._limit_and_offset
@@ -447,17 +451,13 @@ def values_dict(self) -> Dict:
     def comments(self) -> List[str]:
         """
         Return comments from SQL query
-
-        :rtype: List[str]
         """
         return Generalizator(self._raw_query).comments
 
     @property
     def without_comments(self) -> str:
         """
         Removes comments from SQL query
-
-        :rtype: str
         """
         return Generalizator(self._raw_query).without_comments
 
@@ -468,8 +468,6 @@ def generalize(self) -> Optional[str]:
         and replaces them with X or N for numbers.
 
         Based on Mediawiki's DatabaseBase::generalizeSQL
-
-        :rtype: Optional[str]
         """
         return Generalizator(self._raw_query).generalize
 
@@ -489,8 +487,6 @@ def _add_to_columns_subsection(self, keyword: str, column: str):
     def _preprocess_query(self) -> str:
         """
         Perform initial query cleanup
-
-        :rtype str
         """
         if self._raw_query == "":
             return ""

test/test_getting_columns.py

@@ -210,3 +210,18 @@ def test_complex_queries_columns():
         "where": ["cl_type", "cl_to"],
         "order_by": ["cl_sortkey"],
     }
+
+
+def test_columns_and_sql_functions():
+    """
+    See https://github.com/macbre/sql-metadata/issues/125
+    """
+    assert Parser("select max(col3)+avg(col)+1+sum(col2) from dual").columns == [
+        "col3",
+        "col",
+        "col2",
+    ]
+    assert Parser("select avg(col)+sum(col2) from dual").columns == ["col", "col2"]
+    assert Parser(
+        "select count(col)+max(col2)+ min(col3)+ count(distinct  col4) + custom_func(col5) from dual"
+    ).columns == ["col", "col2", "col3", "col4", "col5"]