Merge pull request #54 from macbre/having-query-fix

Add test_query_with_having

Author: macbre

sql_metadata.py

@@ -223,8 +223,11 @@ def get_query_tables(query: str) -> List[str]:
         elif str(token) == '(':
             # reset the last_keyword for INSERT `foo` VALUES(id, bar) ...
             last_keyword = None
-        elif token.is_keyword and str(token) in ['FORCE', 'ORDER']:
-            # reset the last_keyword for "SELECT x FORCE INDEX" queries and "SELECT x ORDER BY"
+        elif token.is_keyword and str(token) in ['FORCE', 'ORDER', 'GROUP BY']:
+            # reset the last_keyword for queries like:
+            # "SELECT x FORCE INDEX"
+            # "SELECT x ORDER BY"
+            # "SELECT x FROM y GROUP BY x"
             last_keyword = None
         elif token.is_keyword and str(token) == 'SELECT' and last_keyword in ['INTO', 'TABLE']:
             # reset the last_keyword for "INSERT INTO SELECT" and "INSERT TABLE SELECT" queries

test/__init__.py

@@ -0,0 +1,5 @@
+"""
+Enable logging when running tests
+"""
+import logging
+logging.basicConfig(level=logging.DEBUG)

test/test_query.py

@@ -343,3 +343,16 @@ def test_sql_server_cte_sales_by_year():
 
     # TODO
     # assert get_query_columns(sales_query) == ['staff_id', 'order_count', 'order_date']
+
+
+def test_table_name_with_group_by():
+    expected_tables = ['SH.sales']
+
+    assert get_query_tables("SELECT s.cust_id,count(s.cust_id) FROM SH.sales s") == expected_tables
+
+    assert get_query_tables("SELECT s.cust_id,count(s.cust_id) FROM SH.sales s GROUP BY s.cust_id") == expected_tables
+
+    assert get_query_tables("""
+SELECT s.cust_id,count(s.cust_id) FROM SH.sales s
+GROUP BY s.cust_id HAVING s.cust_id != '1660' AND s.cust_id != '2'
+    """.strip()) == expected_tables