# select_parser.py
# Copyright 2010,2019 Paul McGuire
#
# a simple SELECT statement parser, taken from SQLite's SELECT statement
# definition at https://www.sqlite.org/lang_select.html
#
from pyparsing import *
ParserElement.enablePackrat()

LPAR,RPAR,COMMA = map(Suppress,"(),")
DOT,STAR = map(Literal, ".*")
select_stmt = Forward().setName("select statement")

# keywords
(UNION, ALL, AND, INTERSECT, EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER,
 CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT, DISTINCT, FROM, WHERE, GROUP, BY,
 HAVING, ORDER, BY, LIMIT, OFFSET, OR) =  map(CaselessKeyword, """UNION, ALL, AND, INTERSECT,
 EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER, CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT,
 DISTINCT, FROM, WHERE, GROUP, BY, HAVING, ORDER, BY, LIMIT, OFFSET, OR""".replace(",","").split())
(CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, ELSE, END, CASE, WHEN, THEN, EXISTS, IN, LIKE, GLOB, REGEXP,
 MATCH, ESCAPE, CURRENT_TIME, CURRENT_DATE, CURRENT_TIMESTAMP) = map(CaselessKeyword,
 """CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, ELSE, END, CASE, WHEN, THEN, EXISTS, IN, LIKE, GLOB, 
 REGEXP, MATCH, ESCAPE, CURRENT_TIME, CURRENT_DATE, CURRENT_TIMESTAMP""".replace(",","").split())
keyword = MatchFirst((UNION, ALL, INTERSECT, EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER,
                      CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT, DISTINCT, FROM, WHERE, GROUP, BY,
                      HAVING, ORDER, BY, LIMIT, OFFSET, CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, ELSE, END,
                      CASE, WHEN, THEN, EXISTS, COLLATE, IN, LIKE, GLOB, REGEXP, MATCH, ESCAPE,
                      CURRENT_TIME, CURRENT_DATE, CURRENT_TIMESTAMP))

identifier = ~keyword + Word(alphas, alphanums+"_")
collation_name = identifier.copy()
column_name = identifier.copy()
column_alias = identifier.copy()
table_name = identifier.copy()
table_alias = identifier.copy()
index_name = identifier.copy()
function_name = identifier.copy()
parameter_name = identifier.copy()
database_name = identifier.copy()

# expression
expr = Forward().setName("expression")

integer = Regex(r"[+-]?\d+")
numeric_literal = Regex(r"\d+(\.\d*)?([eE][+-]?\d+)?")
string_literal = QuotedString("'")
blob_literal = Regex(r"[xX]'[0-9A-Fa-f]+'")
literal_value = (
        numeric_literal
        | string_literal
        | blob_literal
        | NULL
        | CURRENT_TIME
        | CURRENT_DATE
        | CURRENT_TIMESTAMP
        )
bind_parameter = (
    Word("?",nums)
    | Combine(oneOf(": @ $") + parameter_name)
    )
type_name = oneOf("TEXT REAL INTEGER BLOB NULL")

expr_term = (
    CAST + LPAR + expr + AS + type_name + RPAR
    | EXISTS + LPAR + select_stmt + RPAR
    | function_name.setName("function_name") + LPAR + Optional(STAR | delimitedList(expr)) + RPAR
    | literal_value
    | bind_parameter
    | Group(identifier('col_db') + DOT + identifier('col_tab') + DOT + identifier('col'))
    | Group(identifier('col_tab') + DOT + identifier('col'))
    | Group(identifier('col'))
    )

UNARY,BINARY,TERNARY=1,2,3
expr << infixNotation(expr_term,
    [
    (oneOf('- + ~') | NOT, UNARY, opAssoc.RIGHT),
    (ISNULL | NOTNULL | NOT + NULL, UNARY, opAssoc.LEFT),
    ('||', BINARY, opAssoc.LEFT),
    (oneOf('* / %'), BINARY, opAssoc.LEFT),
    (oneOf('+ -'), BINARY, opAssoc.LEFT),
    (oneOf('<< >> & |'), BINARY, opAssoc.LEFT),
    (oneOf('< <= > >='), BINARY, opAssoc.LEFT),
    (oneOf('= == != <>') | IS | IN | LIKE | GLOB | MATCH | REGEXP, BINARY, opAssoc.LEFT),
    ((BETWEEN,AND), TERNARY, opAssoc.LEFT),
    (IN + LPAR + Group(select_stmt | delimitedList(expr)) + RPAR, UNARY, opAssoc.LEFT),
    (AND, BINARY, opAssoc.LEFT),
    (OR, BINARY, opAssoc.LEFT),
    ])

compound_operator = (UNION + Optional(ALL) | INTERSECT | EXCEPT)

ordering_term = Group(expr('order_key')
                      + Optional(COLLATE + collation_name('collate'))
                      + Optional(ASC | DESC)('direction'))

join_constraint = Group(Optional(ON + expr | USING + LPAR + Group(delimitedList(column_name)) + RPAR))

join_op = COMMA | Group(Optional(NATURAL) + Optional(INNER | CROSS | LEFT + OUTER | LEFT | OUTER) + JOIN)

join_source = Forward()
single_source = (
                Group(database_name("database") + DOT + table_name("table*")
                      | table_name("table*"))
                    + Optional(Optional(AS) + table_alias("table_alias*"))
                    + Optional(INDEXED + BY + index_name("name") | NOT + INDEXED)("index")
                | (LPAR + select_stmt + RPAR + Optional(Optional(AS) + table_alias))
                | (LPAR + join_source + RPAR)
                )

join_source <<= (Group(single_source + OneOrMore(join_op + single_source + join_constraint))
                 | single_source)

# result_column = "*" | table_name + "." + "*" | Group(expr + Optional(Optional(AS) + column_alias))
result_column = Group(STAR('col')
                      | table_name('col_table') + DOT + STAR('col')
                      | expr('col') + Optional(Optional(AS) + column_alias('alias'))
                      )

select_core = (SELECT + Optional(DISTINCT | ALL) + Group(delimitedList(result_column))("columns")
               + Optional(FROM + join_source("from*"))
               + Optional(WHERE + expr("where_expr"))
               + Optional(GROUP + BY + Group(delimitedList(ordering_term))("group_by_terms")
                          + Optional(HAVING + expr("having_expr"))))

select_stmt << (select_core + ZeroOrMore(compound_operator + select_core)
                + Optional(ORDER + BY + Group(delimitedList(ordering_term))("order_by_terms"))
                + Optional(LIMIT + (Group(expr + OFFSET + expr) | Group(expr + COMMA + expr) | expr)("limit"))
                )

tests = """\
    select * from xyzzy where z > 100
    select * from xyzzy where z > 100 order by zz
    select * from xyzzy
    select z.* from xyzzy
    select a, b from test_table where 1=1 and b='yes'
    select a, b from test_table where 1=1 and b in (select bb from foo)
    select z.a, b from test_table where 1=1 and b in (select bb from foo)
    select z.a, b from test_table where 1=1 and b in (select bb from foo) order by b,c desc,d
    select z.a, b from test_table left join test2_table where 1=1 and b in (select bb from foo)
    select a, db.table.b as BBB from db.table where 1=1 and BBB='yes'
    select a, db.table.b as BBB from test_table,db.table where 1=1 and BBB='yes'
    select a, db.table.b as BBB from test_table,db.table where 1=1 and BBB='yes' limit 50
    select a, b from test_table where (1=1 or 2=3) and b='yes' group by zx having b=2 order by 1
    SELECT emp.ename as e FROM scott.employee as emp
    SELECT ename as e, fname as f FROM scott.employee as emp
    SELECT emp.eid, fname,lname FROM scott.employee as emp
    SELECT ename, lname, emp.eid FROM scott.employee as emp
    select emp.salary * (1.0 + emp.bonus) as salary_plus_bonus from scott.employee as emp
"""
select_stmt.runTests(tests)