From 617b8f6cd3c55bacf2c80130901508518753f7e1 Mon Sep 17 00:00:00 2001
From: Zi-Xuan Fu <r33s3n6@gmail.com>
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'])