initial support for udf, inline k9 block, assumptions

mo_sql_parsing/keywords.py

@@ -33,6 +33,7 @@ EXCEPT = keyword("except")
 FETCH = keyword("fetch").suppress()
 FROM = keyword("from").suppress()
 FULL = keyword("full")
+FUNCTION = keyword("function").suppress()
 GROUP = keyword("group").suppress()
 HAVING = keyword("having").suppress()
 INNER = keyword("inner")
@@ -109,6 +110,7 @@ INDF = (
     # https://prestodb.io/docs/current/functions/comparison.html#is-distinct-from-and-is-not-distinct-from
     keyword("is not distinct from").set_parser_name("ne!")
 )
+FASSIGN = Literal(":=").set_parser_name("fassign") # Assignment in UDFs
 NEQ = (Literal("!=") | Literal("<>")).set_parser_name("neq")
 LAMBDA = Literal("->").set_parser_name("lambda")
 
@@ -181,6 +183,7 @@ RESERVED = MatchFirst([
     FOREIGN,
     FROM,
     FULL,
+    FUNCTION,
     GROUP_BY,
     GROUP,
     HAVING,
@@ -224,7 +227,10 @@ RESERVED = MatchFirst([
     WITH,
     WITHIN,
 ])
-
+L_INLINE = Literal("<k>").suppress()
+R_INLINE = Literal("</k>").suppress()
+LBRACE = Literal("{").suppress()
+RBRACE = Literal("}").suppress()
 LB = Literal("(").suppress()
 RB = Literal(")").suppress()
 EQ = Char("=").suppress()
@@ -282,6 +288,7 @@ precedence = {
     "lambda": 12,
     "join": 18,
     "list": 18,
+    "function": 30,
     "select": 30,
     "from": 30,
     "window": 35,

mo_sql_parsing/sql_parser.py

@@ -7,7 +7,9 @@
 # Contact: Kyle Lahnakoski (kyle@lahnakoski.com)
 #
 
-
+from operator import add
+from textwrap import indent
+from mo_parsing import whitespaces
 from mo_parsing.helpers import restOfLine
 from mo_parsing.infix import delimited_list
 from mo_parsing.whitespaces import NO_WHITESPACE, Whitespace
@@ -76,7 +78,8 @@ def parser(literal_string, ident, sqlserver=False):
         engine.add_ignore(Literal("/*") + SkipTo("*/", include=True))
 
         var_name = ~RESERVED + ident
-
+        
+        inline_kblock = (L_INLINE + SkipTo(R_INLINE, include=True))("k9")
         # EXPRESSIONS
         expr = Forward()
         column_type, column_definition, column_def_references = get_column_type(
@@ -341,6 +344,22 @@ def parser(literal_string, ident, sqlserver=False):
                 + Group(var_name("name") + AS + over_clause("value"))("join")
             )
         ) / to_join_call
+        
+        fassign = Group(var_name("var") + Suppress(FASSIGN) + expr("expr") + Suppress(";"))("assignment")
+        fassigns = fassign + ZeroOrMore(fassign, Whitespace(white=" \t"))
+
+        fbody = (Optional(fassigns) + expr("ret"))
+
+        udf = (
+            FUNCTION 
+            + var_name("fname") 
+            + LB 
+            + Optional(delimited_list(var_name)("params")) 
+            + RB 
+            + LBRACE
+            + fbody
+            + RBRACE
+        )
 
         selection = (
             (SELECT + DISTINCT + ON + LB)
@@ -407,11 +426,13 @@ def parser(literal_string, ident, sqlserver=False):
             )
             + RB,
         )
-        assumption = (ASSUMING + (ASC|DESC)("assumption"))
+
+        assumption = Group((ASC|DESC) ("ord") + var_name("attrib"))
+        assumptions = (ASSUMING + Group(delimited_list(assumption))("assumptions"))
 
         table_source << Group(
             ((LB + query + RB) | stack | call_function | var_name)("value")
-            + Optional(assumption)
+            + Optional(assumptions)
             + Optional(flag("with ordinality"))
             + Optional(tablesample)
             + alias
@@ -600,7 +621,9 @@ def parser(literal_string, ident, sqlserver=False):
         ) / to_json_call
 
         return (
-            query
+            inline_kblock
+            | udf
+            | query
             | (insert | update | delete)
             | (create_table | create_view | create_cache | create_index)
             | (drop_table | drop_view | drop_index)