gh-79097: Add support for aggregate window functions in sqlite3

Doc/includes/sqlite3/sumintwindow.py

@@ -0,0 +1,46 @@
+# Example taken from https://www.sqlite.org/windowfunctions.html#udfwinfunc
+import sqlite3
+
+
+class WindowSumInt:
+ def __init__(self):
+ self.count = 0
+
+ def step(self, value):
+ """Adds a row to the current window."""
+ self.count += value
+
+ def value(self):
+ """Returns the current value of the aggregate."""
+ return self.count
+
+ def inverse(self, value):
+ """Removes a row from the current window."""
+ self.count -= value
+
+ def finalize(self):
+ """Returns the final value of the aggregate.
+
+ Any clean-up actions should be placed here.
+ """
+ return self.count
+
+
+con = sqlite3.connect(":memory:")
+cur = con.execute("create table test(x, y)")
+values = [
+ ("a", 4),
+ ("b", 5),
+ ("c", 3),
+ ("d", 8),
+ ("e", 1),
+]
+cur.executemany("insert into test values(?, ?)", values)
+con.create_window_function("sumint", 1, WindowSumInt)
+cur.execute("""
+ select x, sumint(y) over (
+ order by x rows between 1 preceding and 1 following
+ ) as sum_y
+ from test order by x
+""")
+print(cur.fetchall())

Doc/library/sqlite3.rst

@@ -473,6 +473,35 @@ Connection Objects
  .. literalinclude:: ../includes/sqlite3/mysumaggr.py
 
 
+ .. method:: create_window_function(name, num_params, aggregate_class, /)
+
+ Creates user-defined aggregate window function *name*.
+
+ *aggregate_class* must implement the following methods:
+
+ * ``step``: adds a row to the current window
+ * ``value``: returns the current value of the aggregate
+ * ``inverse``: removes a row from the current window
+ * ``finalize``: returns the final value of the aggregate
+
+ ``step`` and ``value`` accept *num_params* number of parameters,
+ unless *num_params* is ``-1``, in which case they may take any number of
+ arguments. ``finalize`` and ``value`` can return any of the types
+ supported by SQLite:
+ :class:`bytes`, :class:`str`, :class:`int`, :class:`float`, and
+ :const:`None`. Call :meth:`create_window_function` with
+ *aggregate_class* set to :const:`None` to clear window function *name*.
+
+ Aggregate window functions are supported by SQLite 3.25.0 and higher.
+ :exc:`NotSupportedError` will be raised if used with older versions.
+
+ .. versionadded:: 3.11
+
+ Example:
+
+ .. literalinclude:: ../includes/sqlite3/sumintwindow.py
+
+
  .. method:: create_collation(name, callable)
 
  Creates a collation with the specified *name* and *callable*. The callable will

Doc/whatsnew/3.11.rst

@@ -389,6 +389,10 @@ sqlite3
  serializing and deserializing databases.
  (Contributed by Erlend E. Aasland in :issue:`41930`.)
 
+* Add :meth:`~sqlite3.Connection.create_window_function` to
+ :class:`sqlite3.Connection` for creating aggregate window functions.
+ (Contributed by Erlend E. Aasland in :issue:`34916`.)
+
 
 sys
 ---

Lib/test/test_sqlite3/test_dbapi.py

@@ -1084,6 +1084,8 @@ def test_check_connection_thread(self):
  if hasattr(sqlite.Connection, "serialize"):
  fns.append(lambda: self.con.serialize())
  fns.append(lambda: self.con.deserialize(b""))
+ if sqlite.sqlite_version_info >= (3, 25, 0):
+ fns.append(lambda: self.con.create_window_function("foo", 0, None))
 
  for fn in fns:
  with self.subTest(fn=fn):

Lib/test/test_sqlite3/test_userfunctions.py

@@ -27,9 +27,9 @@
 import re
 import sys
 import unittest
-import unittest.mock
 import sqlite3 as sqlite
 
+from unittest.mock import Mock, patch
 from test.support import bigmemtest, catch_unraisable_exception, gc_collect
 
 from test.test_sqlite3.test_dbapi import cx_limit
@@ -393,7 +393,7 @@ def append_result(arg):
  # indices, which allows testing based on syntax, iso. the query optimizer.
  @unittest.skipIf(sqlite.sqlite_version_info < (3, 8, 3), "Requires SQLite 3.8.3 or higher")
  def test_func_non_deterministic(self):
- mock = unittest.mock.Mock(return_value=None)
+ mock = Mock(return_value=None)
  self.con.create_function("nondeterministic", 0, mock, deterministic=False)
  if sqlite.sqlite_version_info < (3, 15, 0):
  self.con.execute("select nondeterministic() = nondeterministic()")
@@ -404,7 +404,7 @@ def test_func_non_deterministic(self):
 
  @unittest.skipIf(sqlite.sqlite_version_info < (3, 8, 3), "Requires SQLite 3.8.3 or higher")
  def test_func_deterministic(self):
- mock = unittest.mock.Mock(return_value=None)
+ mock = Mock(return_value=None)
  self.con.create_function("deterministic", 0, mock, deterministic=True)
  if sqlite.sqlite_version_info < (3, 15, 0):
  self.con.execute("select deterministic() = deterministic()")
@@ -482,6 +482,164 @@ def test_func_return_illegal_value(self):
  self.con.execute, "select badreturn()")
 
 
+class WindowSumInt:
+ def __init__(self):
+ self.count = 0
+
+ def step(self, value):
+ self.count += value
+
+ def value(self):
+ return self.count
+
+ def inverse(self, value):
+ self.count -= value
+
+ def finalize(self):
+ return self.count
+
+class BadWindow(Exception):
+ pass
+
+
+@unittest.skipIf(sqlite.sqlite_version_info < (3, 25, 0),
+ "Requires SQLite 3.25.0 or newer")
+class WindowFunctionTests(unittest.TestCase):
+ def setUp(self):
+ self.con = sqlite.connect(":memory:")
+ self.cur = self.con.cursor()
+
+ # Test case taken from https://www.sqlite.org/windowfunctions.html#udfwinfunc
+ values = [
+ ("a", 4),
+ ("b", 5),
+ ("c", 3),
+ ("d", 8),
+ ("e", 1),
+ ]
+ with self.con:
+ self.con.execute("create table test(x, y)")
+ self.con.executemany("insert into test values(?, ?)", values)
+ self.expected = [
+ ("a", 9),
+ ("b", 12),
+ ("c", 16),
+ ("d", 12),
+ ("e", 9),
+ ]
+ self.query = """
+ select x, %s(y) over (
+ order by x rows between 1 preceding and 1 following
+ ) as sum_y
+ from test order by x
+ """
+ self.con.create_window_function("sumint", 1, WindowSumInt)
+
+ def test_win_sum_int(self):
+ self.cur.execute(self.query % "sumint")
+ self.assertEqual(self.cur.fetchall(), self.expected)
+
+ def test_win_error_on_create(self):
+ self.assertRaises(sqlite.ProgrammingError,
+ self.con.create_window_function,
+ "shouldfail", -100, WindowSumInt)
+
+ @with_tracebacks(BadWindow)
+ def test_win_exception_in_method(self):
+ for meth in "__init__", "step", "value", "inverse":
+ with self.subTest(meth=meth):
+ with patch.object(WindowSumInt, meth, side_effect=BadWindow):
+ name = f"exc_{meth}"
+ self.con.create_window_function(name, 1, WindowSumInt)
+ msg = f"'{meth}' method raised error"
+ with self.assertRaisesRegex(sqlite.OperationalError, msg):
+ self.cur.execute(self.query % name)
+ self.cur.fetchall()
+
+ @with_tracebacks(BadWindow)
+ def test_win_exception_in_finalize(self):
+ # Note: SQLite does not (as of version 3.38.0) propagate finalize
+ # callback errors to sqlite3_step(); this implies that OperationalError
+ # is _not_ raised.
+ with patch.object(WindowSumInt, "finalize", side_effect=BadWindow):
+ name = f"exception_in_finalize"
+ self.con.create_window_function(name, 1, WindowSumInt)
+ self.cur.execute(self.query % name)
+ self.cur.fetchall()
+
+ @with_tracebacks(AttributeError)
+ def test_win_missing_method(self):
+ class MissingValue:
+ def step(self, x): pass
+ def inverse(self, x): pass
+ def finalize(self): return 42
+
+ class MissingInverse:
+ def step(self, x): pass
+ def value(self): return 42
+ def finalize(self): return 42
+
+ class MissingStep:
+ def value(self): return 42
+ def inverse(self, x): pass
+ def finalize(self): return 42
+
+ dataset = (
+ ("step", MissingStep),
+ ("value", MissingValue),
+ ("inverse", MissingInverse),
+ )
+ for meth, cls in dataset:
+ with self.subTest(meth=meth, cls=cls):
+ name = f"exc_{meth}"
+ self.con.create_window_function(name, 1, cls)
+ with self.assertRaisesRegex(sqlite.OperationalError,
+ f"'{meth}' method not defined"):
+ self.cur.execute(self.query % name)
+ self.cur.fetchall()
+
+ @with_tracebacks(AttributeError)
+ def test_win_missing_finalize(self):
+ # Note: SQLite does not (as of version 3.38.0) propagate finalize
+ # callback errors to sqlite3_step(); this implies that OperationalError
+ # is _not_ raised.
+ class MissingFinalize:
+ def step(self, x): pass
+ def value(self): return 42
+ def inverse(self, x): pass
+
+ name = "missing_finalize"
+ self.con.create_window_function(name, 1, MissingFinalize)
+ self.cur.execute(self.query % name)
+ self.cur.fetchall()
+
+ def test_win_clear_function(self):
+ self.con.create_window_function("sumint", 1, None)
+ self.assertRaises(sqlite.OperationalError, self.cur.execute,
+ self.query % "sumint")
+
+ def test_win_redefine_function(self):
+ # Redefine WindowSumInt; adjust the expected results accordingly.
+ class Redefined(WindowSumInt):
+ def step(self, value): self.count += value * 2
+ def inverse(self, value): self.count -= value * 2
+ expected = [(v[0], v[1]*2) for v in self.expected]
+
+ self.con.create_window_function("sumint", 1, Redefined)
+ self.cur.execute(self.query % "sumint")
+ self.assertEqual(self.cur.fetchall(), expected)
+
+ def test_win_error_value_return(self):
+ class ErrorValueReturn:
+ def __init__(self): pass
+ def step(self, x): pass
+ def value(self): return 1 << 65
+
+ self.con.create_window_function("err_val_ret", 1, ErrorValueReturn)
+ self.assertRaisesRegex(sqlite.DataError, "string or blob too big",
+ self.cur.execute, self.query % "err_val_ret")
+
+
 class AggregateTests(unittest.TestCase):
  def setUp(self):
  self.con = sqlite.connect(":memory:")
@@ -527,10 +685,10 @@ def test_aggr_no_step(self):
 
  def test_aggr_no_finalize(self):
  cur = self.con.cursor()
- with self.assertRaises(sqlite.OperationalError) as cm:
+ msg = "user-defined aggregate's 'finalize' method not defined"
+ with self.assertRaisesRegex(sqlite.OperationalError, msg):
  cur.execute("select nofinalize(t) from test")
  val = cur.fetchone()[0]
- self.assertEqual(str(cm.exception), "user-defined aggregate's 'finalize' method raised error")
 
  @with_tracebacks(ZeroDivisionError, name="AggrExceptionInInit")
  def test_aggr_exception_in_init(self):

Misc/NEWS.d/next/Library/2020-05-24-23-52-03.bpo-40617.lycF9q.rst

@@ -0,0 +1,3 @@
+Add :meth:`~sqlite3.Connection.create_window_function` to
+:class:`sqlite3.Connection` for creating aggregate window functions.
+Patch by Erlend E. Aasland.