Add support for sqlite3 aggregate window functions

Erlend E. Aasland · Erlend E. Aasland · commit bbef5a46f227 · 2021-04-23T11:24:48.000+02:00
See https://docs.python.org/3/library/exceptions.html#NotImplementedError: It should not be used to indicate that an operator or method is not meant to be supported at all – in that case either leave the operator / method undefined or, if a subclass, set it to None.
diff --git a/Doc/includes/sqlite3/sumintwindow.py b/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):
+        """This method is invoked to add a row to the current window."""
+        self.count += value
+
+    def value(self):
+        """This method is invoked to return the current value of the aggregate."""
+        return self.count
+
+    def inverse(self, value):
+        """This method is invoked to remove a row from the current window."""
+        self.count -= value
+
+    def finalize(self):
+        """This method is invoked to return the current 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())
diff --git a/Doc/library/sqlite3.rst b/Doc/library/sqlite3.rst
@@ -379,6 +379,52 @@ Connection Objects
       .. literalinclude:: ../includes/sqlite3/mysumaggr.py
 
 
+   .. method:: create_window_function(name, num_params, aggregate_class, /, *,
+                                      deterministic=False, innocuous=False,
+                                      directonly=False)
+
+      Creates a user-defined aggregate window function. Aggregate window
+      functions are supported by SQLite 3.25.0 and higher.
+      :exc:`NotSupportedError` will be raised if used with older
+      versions.
+
+      The aggregate class must implement ``step`` and ``inverse``
+      methods, which accept the number of parameters *num_params* (if
+      *num_params* is -1, the function may take any number of arguments),
+      and ``finalize`` and ``value`` methods which return the final and
+      the current result of the aggregate.
+
+      The ``finalize`` and ``value`` methods can return any of the types
+      supported by SQLite: :class:`bytes`, :class:`str`, :class:`int`,
+      :class:`float` and :const:`None`.
+
+      If *deterministic* is :const:`True`, the created function is marked as
+      `deterministic <https://sqlite.org/deterministic.html>`_, which
+      allows SQLite to perform additional optimizations.  This flag is
+      supported by SQLite 3.8.3 or higher. :exc:`NotSupportedError` will
+      be raised if used with older versions.
+
+      If *innocuous* is :const:`True`, the created function is marked as
+      `innocuous <https://sqlite.org/deterministic.html>`_, which
+      indicates to SQLite that it is unlikely to cause problems, even if
+      misused.  This flag is supported by SQLite 3.31.0 or higher.
+      :exc:`NotSupportedError` will be raised if used with older
+      versions.
+
+      If *directonly* is :const:`True`, the created function is marked as
+      `directonly <https://sqlite.org/deterministic.html>`_, which
+      means that it may only be invoked from top-level SQL.  This flag
+      is an SQLite security feature that is recommended for all
+      user-defined SQL functions.  This flag is supported by SQLite
+      3.30.0 or higher.  :exc:`NotSupportedError` will be raised if used
+      with older versions.
+
+      .. versionadded:: 3.10
+
+      Example:
+
+      .. literalinclude:: ../includes/sqlite3/sumintwindow.py
+
    .. method:: create_collation(name, callable)
 
       Creates a collation with the specified *name* and *callable*. The callable will
diff --git a/Lib/sqlite3/test/userfunctions.py b/Lib/sqlite3/test/userfunctions.py
@@ -321,6 +321,85 @@ def test_func_deterministic_keyword_only(self):
             self.con.create_function("deterministic", 0, int, True)
 
 
+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
+
+
+@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:")
+
+        # Test case taken from https://www.sqlite.org/windowfunctions.html#udfwinfunc
+        values = [
+            ("a", 4),
+            ("b", 5),
+            ("c", 3),
+            ("d", 8),
+            ("e", 1),
+        ]
+        self.cur = self.con.execute("create table test(x, y)")
+        self.cur.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_sum_int(self):
+        self.cur.execute(self.query % "sumint")
+        self.assertEqual(self.cur.fetchall(), self.expected)
+
+    def test_error_on_create(self):
+        with self.assertRaises(sqlite.OperationalError):
+            self.con.create_window_function("shouldfail", -100, WindowSumInt)
+
+    def test_exception_in_method(self):
+        for meth in ["step", "value", "inverse"]:
+            with unittest.mock.patch.object(WindowSumInt, meth,
+                                            side_effect=Exception):
+                func = f"exc_{meth}"
+                self.con.create_window_function(func, 1, WindowSumInt)
+                with self.assertRaises(sqlite.OperationalError):
+                    self.cur.execute(self.query % func)
+                    ret = self.cur.fetchall()
+
+    def test_clear_function(self):
+        self.con.create_window_function("sumint", 1, None)
+        with self.assertRaises(sqlite.OperationalError):
+            self.cur.execute(self.query % "sumint")
+
+    def test_redefine_function(self):
+        class Redefined(WindowSumInt):
+            pass
+        self.con.create_window_function("sumint", 1, Redefined)
+        self.cur.execute(self.query % "sumint")
+        self.assertEqual(self.cur.fetchall(), self.expected)
+
+
 class AggregateTests(unittest.TestCase):
     def setUp(self):
         self.con = sqlite.connect(":memory:")
@@ -510,6 +589,7 @@ def suite():
         AuthorizerRaiseExceptionTests,
         AuthorizerTests,
         FunctionTests,
+        WindowFunctionTests,
     ]
     return unittest.TestSuite(
         [unittest.TestLoader().loadTestsFromTestCase(t) for t in tests]
diff --git a/Modules/_sqlite/clinic/connection.c.h b/Modules/_sqlite/clinic/connection.c.h
@@ -156,6 +156,109 @@ pysqlite_connection_create_function(pysqlite_Connection *self, PyObject *const *
     return return_value;
 }
 
+#if defined(HAVE_WINDOW_FUNCTIONS)
+
+PyDoc_STRVAR(pysqlite_connection_create_window_function__doc__,
+"create_window_function($self, name, num_params, aggregate_class, /, *,\n"
+"                       deterministic=False, directonly=False,\n"
+"                       innocuous=False)\n"
+"--\n"
+"\n"
+"Creates or redefines an aggregate window function. Non-standard.\n"
+"\n"
+"  name\n"
+"    The name of the SQL aggregate window function to be created or\n"
+"    redefined.\n"
+"  num_params\n"
+"    The number of arguments that the SQL aggregate window function\n"
+"    takes.\n"
+"  aggregate_class\n"
+"    A class with step(), finalize(), value(), and inverse() methods.\n"
+"    Set to None to clear the window function.");
+
+#define PYSQLITE_CONNECTION_CREATE_WINDOW_FUNCTION_METHODDEF    \
+    {"create_window_function", (PyCFunction)(void(*)(void))pysqlite_connection_create_window_function, METH_FASTCALL|METH_KEYWORDS, pysqlite_connection_create_window_function__doc__},
+
+static PyObject *
+pysqlite_connection_create_window_function_impl(pysqlite_Connection *self,
+                                                const char *name,
+                                                int num_params,
+                                                PyObject *aggregate_class,
+                                                int deterministic,
+                                                int directonly,
+                                                int innocuous);
+
+static PyObject *
+pysqlite_connection_create_window_function(pysqlite_Connection *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"", "", "", "deterministic", "directonly", "innocuous", NULL};
+    static _PyArg_Parser _parser = {NULL, _keywords, "create_window_function", 0};
+    PyObject *argsbuf[6];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 3;
+    const char *name;
+    int num_params;
+    PyObject *aggregate_class;
+    int deterministic = 0;
+    int directonly = 0;
+    int innocuous = 0;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 3, 3, 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    if (!PyUnicode_Check(args[0])) {
+        _PyArg_BadArgument("create_window_function", "argument 1", "str", args[0]);
+        goto exit;
+    }
+    Py_ssize_t name_length;
+    name = PyUnicode_AsUTF8AndSize(args[0], &name_length);
+    if (name == NULL) {
+        goto exit;
+    }
+    if (strlen(name) != (size_t)name_length) {
+        PyErr_SetString(PyExc_ValueError, "embedded null character");
+        goto exit;
+    }
+    num_params = _PyLong_AsInt(args[1]);
+    if (num_params == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    aggregate_class = args[2];
+    if (!noptargs) {
+        goto skip_optional_kwonly;
+    }
+    if (args[3]) {
+        deterministic = PyObject_IsTrue(args[3]);
+        if (deterministic < 0) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[4]) {
+        directonly = PyObject_IsTrue(args[4]);
+        if (directonly < 0) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    innocuous = PyObject_IsTrue(args[5]);
+    if (innocuous < 0) {
+        goto exit;
+    }
+skip_optional_kwonly:
+    return_value = pysqlite_connection_create_window_function_impl(self, name, num_params, aggregate_class, deterministic, directonly, innocuous);
+
+exit:
+    return return_value;
+}
+
+#endif /* defined(HAVE_WINDOW_FUNCTIONS) */
+
 PyDoc_STRVAR(pysqlite_connection_create_aggregate__doc__,
 "create_aggregate($self, /, name, n_arg, aggregate_class)\n"
 "--\n"
@@ -703,11 +806,15 @@ pysqlite_connection_exit(pysqlite_Connection *self, PyObject *const *args, Py_ss
     return return_value;
 }
 
+#ifndef PYSQLITE_CONNECTION_CREATE_WINDOW_FUNCTION_METHODDEF
+    #define PYSQLITE_CONNECTION_CREATE_WINDOW_FUNCTION_METHODDEF
+#endif /* !defined(PYSQLITE_CONNECTION_CREATE_WINDOW_FUNCTION_METHODDEF) */
+
 #ifndef PYSQLITE_CONNECTION_ENABLE_LOAD_EXTENSION_METHODDEF
     #define PYSQLITE_CONNECTION_ENABLE_LOAD_EXTENSION_METHODDEF
 #endif /* !defined(PYSQLITE_CONNECTION_ENABLE_LOAD_EXTENSION_METHODDEF) */
 
 #ifndef PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF
     #define PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF
 #endif /* !defined(PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF) */
-/*[clinic end generated code: output=c1bf09db3bcd0105 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=54e136772a234bed input=a9049054013a1b77]*/
diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c
