bpo-45243: Add support for setting/getting sqlite3 connection limits

pablogsal · Nov 1, 2021 · Sep 19, 2021 · Sep 19, 2021 · Sep 19, 2021 · Sep 19, 2021
commit 39d69ce9ca04d18a8265378926ea5a36ca1fa9db
diff --git a/Lib/sqlite3/test/test_dbapi.py b/Lib/sqlite3/test/test_dbapi.py
@@ -166,6 +166,19 @@ def test_module_constants(self):
             "SQLITE_TOOBIG",
             "SQLITE_TRANSACTION",
             "SQLITE_UPDATE",
+            # Run-time limit categories
+            "SQLITE_LIMIT_LENGTH",
+            "SQLITE_LIMIT_SQL_LENGTH",
+            "SQLITE_LIMIT_COLUMN",
+            "SQLITE_LIMIT_EXPR_DEPTH",
+            "SQLITE_LIMIT_COMPOUND_SELECT",
+            "SQLITE_LIMIT_VDBE_OP",
+            "SQLITE_LIMIT_FUNCTION_ARG",
+            "SQLITE_LIMIT_ATTACHED",
+            "SQLITE_LIMIT_LIKE_PATTERN_LENGTH",
+            "SQLITE_LIMIT_VARIABLE_NUMBER",
+            "SQLITE_LIMIT_TRIGGER_DEPTH",
+            "SQLITE_LIMIT_WORKER_THREADS",
         ]
         if sqlite.version_info >= (3, 7, 17):
             consts += ["SQLITE_NOTICE", "SQLITE_WARNING"]
@@ -331,6 +344,21 @@ def test_drop_unused_refs(self):
             cu = self.cx.execute(f"select {n}")
             self.assertEqual(cu.fetchone()[0], n)
 
+    def test_connection_limits(self):
+        param = sqlite.SQLITE_LIMIT_SQL_LENGTH
+        setval = 10
+        ret1 = self.cx.getlimit(param)
+        try:
+            ret2 = self.cx.setlimit(param, setval)
+            self.assertEqual(ret1, ret2)
+            self.assertEqual(self.cx.getlimit(param), setval)
+            msg = "string or blob too big"
+            self.assertRaisesRegex(sqlite.DataError, msg,
+                                   self.cx.execute, "select 1 as '16'")
+        finally:  # restore old limit
+            print("restoring limit")
+            self.cx.setlimit(param, ret1)
+
 
 class UninitialisedConnectionTests(unittest.TestCase):
     def setUp(self):
@@ -766,6 +794,8 @@ def test_check_connection_thread(self):
             lambda: self.con.set_trace_callback(None),
             lambda: self.con.set_authorizer(None),
             lambda: self.con.create_collation("foo", None),
+            lambda: self.con.setlimit(sqlite.SQLITE_LIMIT_LENGTH, -1),
+            lambda: self.con.getlimit(sqlite.SQLITE_LIMIT_LENGTH),
         ]
         for fn in fns:
             with self.subTest(fn=fn):

diff --git a/Modules/_sqlite/clinic/connection.c.h b/Modules/_sqlite/clinic/connection.c.h
@@ -812,11 +812,88 @@ pysqlite_connection_exit(pysqlite_Connection *self, PyObject *const *args, Py_ss
     return return_value;
 }
 
+PyDoc_STRVAR(setlimit__doc__,
+"setlimit($self, limit, value, /)\n"
+"--\n"
+"\n"
+"Set connection run-time limits. Non-standard.\n"
+"\n"
+"  limit\n"
+"    The limit category to be set.\n"
+"  value\n"
+"    The new limit. If the new limit is a negative number, the limit is\n"
+"    unchanged.\n"
+"\n"
+"Attempts to increase a limit above its hard upper bound are silently truncated\n"
+"to the hard upper bound. Regardless of whether or not the limit was changed,\n"
+"the prior value of the limit is returned.");
+
+#define SETLIMIT_METHODDEF    \
+    {"setlimit", (PyCFunction)(void(*)(void))setlimit, METH_FASTCALL, setlimit__doc__},
+
+static PyObject *
+setlimit_impl(pysqlite_Connection *self, int limit, int value);
+
+static PyObject *
+setlimit(pysqlite_Connection *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    int limit;
+    int value;
+
+    if (!_PyArg_CheckPositional("setlimit", nargs, 2, 2)) {
+        goto exit;
+    }
+    limit = _PyLong_AsInt(args[0]);
+    if (limit == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    value = _PyLong_AsInt(args[1]);
+    if (value == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = setlimit_impl(self, limit, value);
+
+exit:
+    return return_value;
+}
+
+PyDoc_STRVAR(getlimit__doc__,
+"getlimit($self, limit, /)\n"
+"--\n"
+"\n"
+"Get connection run-time limits. Non-standard.\n"
+"\n"
+"  limit\n"
+"    The limit category to be queried.");
+
+#define GETLIMIT_METHODDEF    \
+    {"getlimit", (PyCFunction)getlimit, METH_O, getlimit__doc__},
+
+static PyObject *
+getlimit_impl(pysqlite_Connection *self, int limit);
+
+static PyObject *
+getlimit(pysqlite_Connection *self, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    int limit;
+
+    limit = _PyLong_AsInt(arg);
+    if (limit == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = getlimit_impl(self, limit);
+
+exit:
+    return return_value;
+}
+
 #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=5b7268875f33c016 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=85f8b7247801658e input=a9049054013a1b77]*/
@@ -1892,6 +1892,53 @@ pysqlite_connection_exit_impl(pysqlite_Connection *self, PyObject *exc_type,
     Py_RETURN_FALSE;
 }
 
+/*[clinic input]
+_sqlite3.Connection.setlimit as setlimit
+
+    limit: int
+        The limit category to be set.
+    value: int
+        The new limit. If the new limit is a negative number, the limit is
+        unchanged.
+    /
+
+Set connection run-time limits. Non-standard.
+
+Attempts to increase a limit above its hard upper bound are silently truncated
+to the hard upper bound. Regardless of whether or not the limit was changed,
+the prior value of the limit is returned.
+[clinic start generated code]*/
+
+static PyObject *
+setlimit_impl(pysqlite_Connection *self, int limit, int value)
+/*[clinic end generated code: output=18d15e4be8a7a4ec input=e0990e0c90e747ee]*/
+{
+    if (!pysqlite_check_thread(self) || !pysqlite_check_connection(self)) {
+        return NULL;
+    }
+
+    int previous_value = sqlite3_limit(self->db, limit, value);
+    return PyLong_FromLong(previous_value);
+}
+
+/*[clinic input]
+_sqlite3.Connection.getlimit as getlimit
+
+    limit: int
+        The limit category to be queried.
+    /
+
+Get connection run-time limits. Non-standard.
+[clinic start generated code]*/
+
+static PyObject *
+getlimit_impl(pysqlite_Connection *self, int limit)
+/*[clinic end generated code: output=f90c9c8399170536 input=c339eaae3afe367a]*/
+{
+    return setlimit_impl(self, limit, -1);
+}
+
+
 static const char connection_doc[] =
 PyDoc_STR("SQLite database connection object.");
 
@@ -1923,6 +1970,8 @@ static PyMethodDef connection_methods[] = {
     PYSQLITE_CONNECTION_SET_AUTHORIZER_METHODDEF
     PYSQLITE_CONNECTION_SET_PROGRESS_HANDLER_METHODDEF
     PYSQLITE_CONNECTION_SET_TRACE_CALLBACK_METHODDEF
+    SETLIMIT_METHODDEF
+    GETLIMIT_METHODDEF
     {NULL, NULL}
 };
 

@@ -392,6 +392,19 @@ static int add_integer_constants(PyObject *module) {
 #if SQLITE_VERSION_NUMBER >= 3008003
     ret += PyModule_AddIntMacro(module, SQLITE_RECURSIVE);
 #endif
+    // Run-time limit categories
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_LENGTH);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_SQL_LENGTH);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_COLUMN);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_EXPR_DEPTH);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_COMPOUND_SELECT);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_VDBE_OP);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_FUNCTION_ARG);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_ATTACHED);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_LIKE_PATTERN_LENGTH);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_VARIABLE_NUMBER);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_TRIGGER_DEPTH);
+    ret += PyModule_AddIntMacro(module, SQLITE_LIMIT_WORKER_THREADS);
     return ret;
 }