support hive create function (#500)

* support hive create function:
Reference: https://cwiki.apache.org/confluence/display/Hive/LanguageManual+DDL#LanguageManualDDL-CreateFunction

* Pylint

* func: get_switch_by_create_query set to private method

Author: WangYuhang-CN

sql_metadata/keywords_lists.py

@@ -108,6 +108,7 @@ class TokenType(str, Enum):
     "CREATETABLE": QueryType.CREATE,
     "ALTERTABLE": QueryType.ALTER,
     "DROPTABLE": QueryType.DROP,
+    "CREATEFUNCTION": QueryType.CREATE,
 }
 
 # all the keywords we care for - rest is ignored in assigning

sql_metadata/parser.py

@@ -114,7 +114,9 @@ def query_type(self) -> str:
             )
             .position
         )
-        if tokens[index].normalized in ["CREATE", "ALTER", "DROP"]:
+        if tokens[index].normalized == "CREATE":
+            switch = self._get_switch_by_create_query(tokens, index)
+        elif tokens[index].normalized in ("ALTER", "DROP"):
             switch = tokens[index].normalized + tokens[index + 1].normalized
         else:
             switch = tokens[index].normalized
@@ -1079,3 +1081,19 @@ def _flatten_sqlparse(self):
                             yield tok
             else:
                 yield token
+
+    @staticmethod
+    def _get_switch_by_create_query(tokens: List[SQLToken], index: int) -> str:
+        """
+        Return the switch that creates query type.
+        """
+        switch = tokens[index].normalized + tokens[index + 1].normalized
+
+        # Hive CREATE FUNCTION
+        if any(
+            index + i < len(tokens) and tokens[index + i].normalized == "FUNCTION"
+            for i in (1, 2)
+        ):
+            switch = "CREATEFUNCTION"
+
+        return switch

test/test_query_type.py

@@ -93,3 +93,22 @@ def test_multiple_redundant_parentheses_create():
     """
     parser = Parser(query)
     assert parser.query_type == QueryType.CREATE
+
+
+def test_hive_create_function():
+    query = """
+        CREATE FUNCTION simple_udf AS 'com.example.hive.udf.SimpleUDF' 
+        USING JAR 'hdfs:///user/hive/udfs/simple-udf.jar'
+        WITH SERDEPROPERTIES (
+          "hive.udf.param1"="value1",
+          "hive.udf.param2"="value2"
+        );
+    """
+    parser = Parser(query)
+    assert parser.query_type == QueryType.CREATE
+
+    query = """
+        CREATE TEMPORARY FUNCTION myudf AS 'com.udf.myudf';
+    """
+    parser = Parser(query)
+    assert parser.query_type == QueryType.CREATE