Merge pull request #164 from macbre/handle-leading-comments

Handle leading comments in queries

sql_metadata/parser.py

@@ -1,6 +1,7 @@
 """
 This module provides SQL query parsing functions
 """
+import logging
 import re
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -29,6 +30,8 @@ class Parser:  # pylint: disable=R0902
     """
 
     def __init__(self, sql: str = "") -> None:
+        self._logger = logging.getLogger(self.__class__.__name__)
+
         self._raw_query = sql
         self._query = self._preprocess_query()
         self._query_type = None
@@ -85,12 +88,22 @@ def query_type(self) -> str:
             return self._query_type
         if not self._tokens:
             _ = self.tokens
-        if self._tokens[0].normalized in ["CREATE", "ALTER"]:
-            switch = self._tokens[0].normalized + self._tokens[1].normalized
+
+        # remove comment tokens to not confuse the logic below (see #163)
+        tokens: List[SQLToken] = list(
+            filter(lambda token: not token.is_comment, self._tokens or [])
+        )
+
+        if not tokens:
+            raise ValueError("Empty queries are not supported!")
+
+        if tokens[0].normalized in ["CREATE", "ALTER"]:
+            switch = tokens[0].normalized + tokens[1].normalized
         else:
-            switch = self._tokens[0].normalized
+            switch = tokens[0].normalized
         self._query_type = SUPPORTED_QUERY_TYPES.get(switch, "UNSUPPORTED")
         if self._query_type == "UNSUPPORTED":
+            self._logger.error("Not supported query type: %s", self._raw_query)
             raise ValueError("Not supported query type!")
         return self._query_type
 

test/test_getting_tables.py

@@ -444,3 +444,15 @@ def test_get_tables_with_leading_digits():
     assert ["0020_big_table"] == Parser(
         "SELECT t.val as value, count(*) FROM 0020_big_table"
     ).tables
+
+
+def test_insert_ignore_with_comments():
+    queries = [
+        "INSERT IGNORE /* foo */ INTO bar VALUES (1, '123', '2017-01-01');",
+        "/* foo */ INSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
+        "-- foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
+        "# foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');",
+    ]
+
+    for query in queries:
+        assert ["bar"] == Parser(query).tables

test/test_query_type.py

@@ -0,0 +1,50 @@
+import pytest
+
+from sql_metadata import Parser
+
+
+def test_insert_query():
+    queries = [
+        "INSERT IGNORE /* foo */ INTO bar VALUES (1, '123', '2017-01-01');",
+        "/* foo */ INSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
+        "-- foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
+        "# foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');",
+    ]
+
+    for query in queries:
+        assert "INSERT" == Parser(query).query_type
+
+
+def test_select_query():
+    queries = [
+        "SELECT /* foo */ foo FROM bar",
+        "/* foo */ SELECT foo FROM bar"
+        "-- foo\nSELECT foo FROM bar"
+        "# foo\nSELECT foo FROM bar",
+    ]
+
+    for query in queries:
+        assert "SELECT" == Parser(query).query_type
+
+
+def test_unsupported_query():
+    queries = [
+        "FOO BAR",
+        "DO SOMETHING",
+    ]
+
+    for query in queries:
+        with pytest.raises(ValueError) as ex:
+            _ = Parser(query).query_type
+
+        assert "Not supported query type!" in str(ex.value)
+
+
+def test_empty_query():
+    queries = ["", "/* empty query */"]
+
+    for query in queries:
+        with pytest.raises(ValueError) as ex:
+            _ = Parser(query).query_type
+
+        assert "Empty queries are not supported!" in str(ex.value)

test/test_values.py

@@ -24,8 +24,9 @@ def test_getting_values():
     }
 
     parser = Parser(
-        "INSERT IGNORE INTO `0070_insert_ignore_table` VALUES (9, 2.15, '123', '2017-01-01');"
+        "/* method */ INSERT IGNORE INTO `0070_insert_ignore_table` VALUES (9, 2.15, '123', '2017-01-01');"
     )
+    assert parser.query_type == "INSERT"
     assert parser.values == [9, 2.15, "123", "2017-01-01"]
     assert parser.values_dict == {
         "column_1": 9,