From 617b8f6cd3c55bacf2c80130901508518753f7e1 Mon Sep 17 00:00:00 2001 From: Zi-Xuan Fu Date: Tue, 26 Mar 2024 21:31:51 +0800 Subject: [PATCH] Add OVER clause, and group it into Function (fixes #701) --- sqlparse/engine/grouping.py | 18 +++++++++++++++++- sqlparse/sql.py | 5 +++++ tests/test_grouping.py | 14 ++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py index 9190797a..926a3c1b 100644 --- a/sqlparse/engine/grouping.py +++ b/sqlparse/engine/grouping.py @@ -235,6 +235,16 @@ def group_identifier(tlist): tidx, token = tlist.token_next_by(t=ttypes, idx=tidx) +@recurse(sql.Over) +def group_over(tlist): + tidx, token = tlist.token_next_by(m=sql.Over.M_OPEN) + while token: + nidx, next_ = tlist.token_next(tidx) + if imt(next_, i=sql.Parenthesis, t=T.Name): + tlist.group_tokens(sql.Over, tidx, nidx) + tidx, token = tlist.token_next_by(m=sql.Over.M_OPEN, idx=tidx) + + def group_arrays(tlist): sqlcls = sql.SquareBrackets, sql.Identifier, sql.Function ttypes = T.Name, T.String.Symbol @@ -361,7 +371,12 @@ def group_functions(tlist): while token: nidx, next_ = tlist.token_next(tidx) if isinstance(next_, sql.Parenthesis): - tlist.group_tokens(sql.Function, tidx, nidx) + over_idx, over = tlist.token_next(nidx) + if over and isinstance(over, sql.Over): + eidx = over_idx + else: + eidx = nidx + tlist.group_tokens(sql.Function, tidx, eidx) tidx, token = tlist.token_next_by(t=T.Name, idx=tidx) @@ -412,6 +427,7 @@ def group(stmt): group_for, group_begin, + group_over, group_functions, group_where, group_period, diff --git a/sqlparse/sql.py b/sqlparse/sql.py index 41606dd8..def06797 100644 --- a/sqlparse/sql.py +++ b/sqlparse/sql.py @@ -554,6 +554,11 @@ class Where(TokenList): 'HAVING', 'RETURNING', 'INTO') +class Over(TokenList): + """An OVER clause.""" + M_OPEN = T.Keyword, 'OVER' + + class Having(TokenList): """A HAVING clause.""" M_OPEN = T.Keyword, 'HAVING' diff --git a/tests/test_grouping.py b/tests/test_grouping.py index e90243b5..0bf10c38 100644 --- a/tests/test_grouping.py +++ b/tests/test_grouping.py @@ -185,6 +185,20 @@ def test_grouping_identifier_function(): assert isinstance(p.tokens[0], sql.Identifier) assert isinstance(p.tokens[0].tokens[0], sql.Operation) assert isinstance(p.tokens[0].tokens[0].tokens[0], sql.Function) + p = sqlparse.parse('foo(c1) over win1 as bar')[0] + assert isinstance(p.tokens[0], sql.Identifier) + assert isinstance(p.tokens[0].tokens[0], sql.Function) + assert len(p.tokens[0].tokens[0].tokens) == 4 + assert isinstance(p.tokens[0].tokens[0].tokens[3], sql.Over) + assert isinstance(p.tokens[0].tokens[0].tokens[3].tokens[2], + sql.Identifier) + p = sqlparse.parse('foo(c1) over (partition by c2 order by c3) as bar')[0] + assert isinstance(p.tokens[0], sql.Identifier) + assert isinstance(p.tokens[0].tokens[0], sql.Function) + assert len(p.tokens[0].tokens[0].tokens) == 4 + assert isinstance(p.tokens[0].tokens[0].tokens[3], sql.Over) + assert isinstance(p.tokens[0].tokens[0].tokens[3].tokens[2], + sql.Parenthesis) @pytest.mark.parametrize('s', ['foo+100', 'foo + 100', 'foo*100'])