Support range specification in window function

Fixes: #3049
Co-authored-by: Mike Bayer <mike_mp@zzzcomputing.com>
Change-Id: Ie572095c3e25f70a1e72e1af6858e5edd89fd25e
Pull-request: https://github.com/zzzeek/sqlalchemy/pull/264

3 files changed, 142 insertions, 9 deletions

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 144f2aa47..6d9ab9039 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -810,16 +810,41 @@ class SQLCompiler(Compiled):
             (cast.clause._compiler_dispatch(self, **kwargs),
              cast.typeclause._compiler_dispatch(self, **kwargs))
 
+    def _format_frame_clause(self, range_, **kw):
+        return '%s AND %s' % (
+            "UNBOUNDED PRECEDING"
+            if range_[0] is elements.RANGE_UNBOUNDED
+            else "CURRENT ROW" if range_[0] is elements.RANGE_CURRENT
+            else "%s PRECEDING" % (self.process(range_[0], **kw), ),
+
+            "UNBOUNDED FOLLOWING"
+            if range_[1] is elements.RANGE_UNBOUNDED
+            else "CURRENT ROW" if range_[1] is elements.RANGE_CURRENT
+            else "%s FOLLOWING" % (self.process(range_[1], **kw), )
+        )
+
     def visit_over(self, over, **kwargs):
+        if over.range_:
+            range_ = "RANGE BETWEEN %s" % self._format_frame_clause(
+                over.range_, **kwargs)
+        elif over.rows:
+            range_ = "ROWS BETWEEN %s" % self._format_frame_clause(
+                over.rows, **kwargs)
+        else:
+            range_ = None
+
         return "%s OVER (%s)" % (
             over.element._compiler_dispatch(self, **kwargs),
-            ' '.join(
-                '%s BY %s' % (word, clause._compiler_dispatch(self, **kwargs))
+            ' '.join([
+                '%s BY %s' % (
+                    word, clause._compiler_dispatch(self, **kwargs)
+                )
                 for word, clause in (
                     ('PARTITION', over.partition_by),
                     ('ORDER', over.order_by)
                 )
                 if clause is not None and len(clause)
+            ] + ([range_] if range_ else [])
             )
         )
 
diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py
index e0367f967..e277b28a4 100644
--- a/lib/sqlalchemy/sql/elements.py
+++ b/lib/sqlalchemy/sql/elements.py
@@ -3058,6 +3058,10 @@ class Grouping(ColumnElement):
             self.element.compare(other.element)
 
 
+RANGE_UNBOUNDED = util.symbol("RANGE_UNBOUNDED")
+RANGE_CURRENT = util.symbol("RANGE_CURRENT")
+
+
 class Over(ColumnElement):
     """Represent an OVER clause.
 
@@ -3073,7 +3077,9 @@ class Over(ColumnElement):
     order_by = None
     partition_by = None
 
-    def __init__(self, element, partition_by=None, order_by=None):
+    def __init__(
+            self, element, partition_by=None,
+            order_by=None, range_=None, rows=None):
         """Produce an :class:`.Over` object against a function.
 
         Used against aggregate or so-called "window" functions,
@@ -3082,9 +3088,41 @@ class Over(ColumnElement):
         :func:`~.expression.over` is usually called using
         the :meth:`.FunctionElement.over` method, e.g.::
 
-            func.row_number().over(order_by='x')
+            func.row_number().over(order_by=mytable.c.some_column)
+
+        Would produce::
+
+            ROW_NUMBER() OVER(ORDER BY some_column)
+
+        Ranges are also possible using the :paramref:`.expression.over.range_`
+        and :paramref:`.expression.over.rows` parameters.  These
+        mutually-exclusive parameters each accept a 2-tuple, which contains
+        a combination of integers and None::
+
+            func.row_number().over(order_by=my_table.c.some_column, range_=(None, 0))
+
+        The above would produce::
+
+            ROW_NUMBER() OVER(ORDER BY some_column RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)
+
+        A value of None indicates "unbounded", a
+        value of zero indicates "current row", and negative / positive
+        integers indicate "preceding" and "following":
+
+        * RANGE BETWEEN 5 PRECEDING AND 10 FOLLOWING::
+
+            func.row_number().over(order_by='x', range_=(-5, 10))
+
+        * ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW::
+
+            func.row_number().over(order_by='x', rows=(None, 0))
+
+        * RANGE BETWEEN 2 PRECEDING AND UNBOUNDED FOLLOWING::
+
+            func.row_number().over(order_by='x', range_=(-2, None))
+
+        .. versionadded:: 1.1 support for RANGE / ROWS within a window
 
-        Would produce ``ROW_NUMBER() OVER(ORDER BY x)``.
 
         :param element: a :class:`.FunctionElement`, :class:`.WithinGroup`,
          or other compatible construct.
@@ -3094,12 +3132,21 @@ class Over(ColumnElement):
         :param order_by: a column element or string, or a list
          of such, that will be used as the ORDER BY clause
          of the OVER construct.
+        :param range_: optional range clause for the window.  This is a
+         tuple value which can contain integer values or None, and will
+         render a RANGE BETWEEN PRECEDING / FOLLOWING clause
+
+         .. versionadded:: 1.1
+
+        :param rows: optional rows clause for the window.  This is a tuple
+         value which can contain integer values or None, and will render
+         a ROWS BETWEEN PRECEDING / FOLLOWING clause.
+
+         .. versionadded:: 1.1
 
         This function is also available from the :data:`~.expression.func`
         construct itself via the :meth:`.FunctionElement.over` method.
 
-        .. versionadded:: 0.7
-
         .. seealso::
 
             :data:`.expression.func`
@@ -3117,6 +3164,61 @@ class Over(ColumnElement):
                 *util.to_list(partition_by),
                 _literal_as_text=_literal_as_label_reference)
 
+        if range_:
+            self.range_ = self._interpret_range(range_)
+            if rows:
+                raise exc.ArgumentError(
+                    "'range_' and 'rows' are mutually exclusive")
+            else:
+                self.rows = None
+        elif rows:
+            self.rows = self._interpret_range(rows)
+            self.range_ = None
+        else:
+            self.rows = self.range_ = None
+
+    def _interpret_range(self, range_):
+        if not isinstance(range_, tuple) or len(range_) != 2:
+            raise exc.ArgumentError("2-tuple expected for range/rows")
+
+        if range_[0] is None:
+            preceding = RANGE_UNBOUNDED
+        else:
+            try:
+                preceding = int(range_[0])
+            except ValueError:
+                raise exc.ArgumentError(
+                    "Integer or None expected for preceding value")
+            else:
+                if preceding > 0:
+                    raise exc.ArgumentError(
+                        "Preceding value must be a "
+                        "negative integer, zero, or None")
+                elif preceding < 0:
+                    preceding = literal(abs(preceding))
+                else:
+                    preceding = RANGE_CURRENT
+
+        if range_[1] is None:
+            following = RANGE_UNBOUNDED
+        else:
+            try:
+                following = int(range_[1])
+            except ValueError:
+                raise exc.ArgumentError(
+                    "Integer or None expected for following value")
+            else:
+                if following < 0:
+                    raise exc.ArgumentError(
+                        "Following value must be a positive "
+                        "integer, zero, or None")
+                elif following > 0:
+                    following = literal(following)
+                else:
+                    following = RANGE_CURRENT
+
+        return preceding, following
+
     @property
     def func(self):
         """the element referred to by this :class:`.Over`
diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
index e369f5a61..5c977cd50 100644
--- a/lib/sqlalchemy/sql/functions.py
+++ b/lib/sqlalchemy/sql/functions.py
@@ -94,7 +94,7 @@ class FunctionElement(Executable, ColumnElement, FromClause):
         """
         return self.clause_expr.element
 
-    def over(self, partition_by=None, order_by=None):
+    def over(self, partition_by=None, order_by=None, rows=None, range_=None):
         """Produce an OVER clause against this function.
 
         Used against aggregate or so-called "window" functions,
@@ -114,7 +114,13 @@ class FunctionElement(Executable, ColumnElement, FromClause):
         .. versionadded:: 0.7
 
         """
-        return Over(self, partition_by=partition_by, order_by=order_by)
+        return Over(
+            self,
+            partition_by=partition_by,
+            order_by=order_by,
+            rows=rows,
+            range_=range_
+        )
 
     def within_group(self, *order_by):
         """Produce a WITHIN GROUP (ORDER BY expr) clause against this function.