bpo-45512: Simplify isolation_level handling in `sqlite3` (GH-29053)

3 files changed, 84 insertions, 63 deletions

diff --git a/Modules/_sqlite/clinic/connection.c.h b/Modules/_sqlite/clinic/connection.c.h
index 4e5ee50..5bfc589 100644
--- a/Modules/_sqlite/clinic/connection.c.h
+++ b/Modules/_sqlite/clinic/connection.c.h
@@ -5,7 +5,7 @@ preserve
 static int
 pysqlite_connection_init_impl(pysqlite_Connection *self,
                               const char *database, double timeout,
-                              int detect_types, PyObject *isolation_level,
+                              int detect_types, const char *isolation_level,
                               int check_same_thread, PyObject *factory,
                               int cached_statements, int uri);
 
@@ -22,7 +22,7 @@ pysqlite_connection_init(PyObject *self, PyObject *args, PyObject *kwargs)
     const char *database = NULL;
     double timeout = 5.0;
     int detect_types = 0;
-    PyObject *isolation_level = NULL;
+    const char *isolation_level = "";
     int check_same_thread = 1;
     PyObject *factory = (PyObject*)clinic_state()->ConnectionType;
     int cached_statements = 128;
@@ -63,7 +63,24 @@ pysqlite_connection_init(PyObject *self, PyObject *args, PyObject *kwargs)
         }
     }
     if (fastargs[3]) {
-        isolation_level = fastargs[3];
+        if (fastargs[3] == Py_None) {
+            isolation_level = NULL;
+        }
+        else if (PyUnicode_Check(fastargs[3])) {
+            Py_ssize_t isolation_level_length;
+            isolation_level = PyUnicode_AsUTF8AndSize(fastargs[3], &isolation_level_length);
+            if (isolation_level == NULL) {
+                goto exit;
+            }
+            if (strlen(isolation_level) != (size_t)isolation_level_length) {
+                PyErr_SetString(PyExc_ValueError, "embedded null character");
+                goto exit;
+            }
+        }
+        else {
+            _PyArg_BadArgument("Connection", "argument 'isolation_level'", "str or None", fastargs[3]);
+            goto exit;
+        }
         if (!--noptargs) {
             goto skip_optional_pos;
         }
@@ -834,4 +851,4 @@ exit:
 #ifndef PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF
     #define PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF
 #endif /* !defined(PYSQLITE_CONNECTION_LOAD_EXTENSION_METHODDEF) */
-/*[clinic end generated code: output=d71bf16bef67878f input=a9049054013a1b77]*/
+/*[clinic end generated code: output=663b1e9e71128f19 input=a9049054013a1b77]*/
diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c
index 964673a..e4f0013 100644
--- a/Modules/_sqlite/connection.c
+++ b/Modules/_sqlite/connection.c
@@ -79,7 +79,6 @@ static const char * const begin_statements[] = {
     NULL
 };
 
-static int pysqlite_connection_set_isolation_level(pysqlite_Connection* self, PyObject* isolation_level, void *Py_UNUSED(ignored));
 static void _pysqlite_drop_unused_cursor_references(pysqlite_Connection* self);
 static void free_callback_context(callback_context *ctx);
 static void set_callback_context(callback_context **ctx_pp,
@@ -107,6 +106,30 @@ new_statement_cache(pysqlite_Connection *self, int maxsize)
     return res;
 }
 
+static inline const char *
+begin_stmt_to_isolation_level(const char *begin_stmt)
+{
+    assert(begin_stmt != NULL);
+
+    // All begin statements start with "BEGIN "; add strlen("BEGIN ") to get
+    // the isolation level.
+    return begin_stmt + 6;
+}
+
+static const char *
+get_begin_statement(const char *level)
+{
+    assert(level != NULL);
+    for (int i = 0; begin_statements[i] != NULL; i++) {
+        const char *stmt = begin_statements[i];
+        const char *candidate = begin_stmt_to_isolation_level(stmt);
+        if (sqlite3_stricmp(level, candidate) == 0) {
+            return begin_statements[i];
+        }
+    }
+    return NULL;
+}
+
 /*[python input]
 class FSConverter_converter(CConverter):
     type = "const char *"
@@ -124,7 +147,7 @@ _sqlite3.Connection.__init__ as pysqlite_connection_init
     database: FSConverter
     timeout: double = 5.0
     detect_types: int = 0
-    isolation_level: object = NULL
+    isolation_level: str(accept={str, NoneType}) = ""
     check_same_thread: bool(accept={int}) = True
     factory: object(c_default='(PyObject*)clinic_state()->ConnectionType') = ConnectionType
     cached_statements: int = 128
@@ -134,10 +157,10 @@ _sqlite3.Connection.__init__ as pysqlite_connection_init
 static int
 pysqlite_connection_init_impl(pysqlite_Connection *self,
                               const char *database, double timeout,
-                              int detect_types, PyObject *isolation_level,
+                              int detect_types, const char *isolation_level,
                               int check_same_thread, PyObject *factory,
                               int cached_statements, int uri)
-/*[clinic end generated code: output=bc39e55eb0b68783 input=f8d1f7efc0d84104]*/
+/*[clinic end generated code: output=d8c37afc46d318b0 input=adfb29ac461f9e61]*/
 {
     int rc;
 
@@ -148,8 +171,6 @@ pysqlite_connection_init_impl(pysqlite_Connection *self,
     pysqlite_state *state = pysqlite_get_state_by_type(Py_TYPE(self));
     self->state = state;
 
-    self->begin_statement = NULL;
-
     Py_CLEAR(self->statement_cache);
     Py_CLEAR(self->cursors);
 
@@ -174,20 +195,16 @@ pysqlite_connection_init_impl(pysqlite_Connection *self,
         return -1;
     }
 
-    if (!isolation_level) {
-        isolation_level = PyUnicode_FromString("");
-        if (!isolation_level) {
+    if (isolation_level) {
+        const char *stmt = get_begin_statement(isolation_level);
+        if (stmt == NULL) {
             return -1;
         }
-    } else {
-        Py_INCREF(isolation_level);
+        self->begin_statement = stmt;
     }
-    Py_CLEAR(self->isolation_level);
-    if (pysqlite_connection_set_isolation_level(self, isolation_level, NULL) != 0) {
-        Py_DECREF(isolation_level);
-        return -1;
+    else {
+        self->begin_statement = NULL;
     }
-    Py_DECREF(isolation_level);
 
     self->statement_cache = new_statement_cache(self, cached_statements);
     if (self->statement_cache == NULL) {
@@ -268,7 +285,6 @@ static int
 connection_traverse(pysqlite_Connection *self, visitproc visit, void *arg)
 {
     Py_VISIT(Py_TYPE(self));
-    Py_VISIT(self->isolation_level);
     Py_VISIT(self->statement_cache);
     Py_VISIT(self->cursors);
     Py_VISIT(self->row_factory);
@@ -292,7 +308,6 @@ clear_callback_context(callback_context *ctx)
 static int
 connection_clear(pysqlite_Connection *self)
 {
-    Py_CLEAR(self->isolation_level);
     Py_CLEAR(self->statement_cache);
     Py_CLEAR(self->cursors);
     Py_CLEAR(self->row_factory);
@@ -1317,7 +1332,12 @@ static PyObject* pysqlite_connection_get_isolation_level(pysqlite_Connection* se
     if (!pysqlite_check_connection(self)) {
         return NULL;
     }
-    return Py_NewRef(self->isolation_level);
+    if (self->begin_statement != NULL) {
+        const char *stmt = self->begin_statement;
+        const char *iso_level = begin_stmt_to_isolation_level(stmt);
+        return PyUnicode_FromString(iso_level);
+    }
+    Py_RETURN_NONE;
 }
 
 static PyObject* pysqlite_connection_get_total_changes(pysqlite_Connection* self, void* unused)
@@ -1347,53 +1367,40 @@ pysqlite_connection_set_isolation_level(pysqlite_Connection* self, PyObject* iso
         PyErr_SetString(PyExc_AttributeError, "cannot delete attribute");
         return -1;
     }
-    if (isolation_level == Py_None) {
-        /* We might get called during connection init, so we cannot use
-         * pysqlite_connection_commit() here. */
-        if (self->db && !sqlite3_get_autocommit(self->db)) {
-            int rc;
-            Py_BEGIN_ALLOW_THREADS
-            rc = sqlite3_exec(self->db, "COMMIT", NULL, NULL, NULL);
-            Py_END_ALLOW_THREADS
-            if (rc != SQLITE_OK) {
-                return _pysqlite_seterror(self->state, self->db);
-            }
-        }
-
+    if (Py_IsNone(isolation_level)) {
         self->begin_statement = NULL;
-    } else {
-        const char * const *candidate;
-        PyObject *uppercase_level;
-        _Py_IDENTIFIER(upper);
-
-        if (!PyUnicode_Check(isolation_level)) {
-            PyErr_Format(PyExc_TypeError,
-                         "isolation_level must be a string or None, not %.100s",
-                         Py_TYPE(isolation_level)->tp_name);
+
+        // Execute a COMMIT to re-enable autocommit mode
+        PyObject *res = pysqlite_connection_commit_impl(self);
+        if (res == NULL) {
             return -1;
         }
-
-        uppercase_level = _PyObject_CallMethodIdOneArg(
-                        (PyObject *)&PyUnicode_Type, &PyId_upper,
-                        isolation_level);
-        if (!uppercase_level) {
+        Py_DECREF(res);
+    }
+    else if (PyUnicode_Check(isolation_level)) {
+        Py_ssize_t len;
+        const char *cstr_level = PyUnicode_AsUTF8AndSize(isolation_level, &len);
+        if (cstr_level == NULL) {
             return -1;
         }
-        for (candidate = begin_statements; *candidate; candidate++) {
-            if (_PyUnicode_EqualToASCIIString(uppercase_level, *candidate + 6))
-                break;
+        if (strlen(cstr_level) != (size_t)len) {
+            PyErr_SetString(PyExc_ValueError, "embedded null character");
+            return -1;
         }
-        Py_DECREF(uppercase_level);
-        if (!*candidate) {
+        const char *stmt = get_begin_statement(cstr_level);
+        if (stmt == NULL) {
             PyErr_SetString(PyExc_ValueError,
-                            "invalid value for isolation_level");
+                            "isolation_level string must be '', 'DEFERRED', "
+                            "'IMMEDIATE', or 'EXCLUSIVE'");
             return -1;
         }
-        self->begin_statement = *candidate;
+        self->begin_statement = stmt;
+    }
+    else {
+        PyErr_SetString(PyExc_TypeError,
+                        "isolation_level must be str or None");
+        return -1;
     }
-
-    Py_INCREF(isolation_level);
-    Py_XSETREF(self->isolation_level, isolation_level);
     return 0;
 }
 
diff --git a/Modules/_sqlite/connection.h b/Modules/_sqlite/connection.h
index 6a2aa1c..7baf63f 100644
--- a/Modules/_sqlite/connection.h
+++ b/Modules/_sqlite/connection.h
@@ -49,9 +49,6 @@ typedef struct
      * bitwise combination thereof makes sense */
     int detect_types;
 
-    /* None for autocommit, otherwise a PyUnicode with the isolation level */
-    PyObject* isolation_level;
-
     /* NULL for autocommit, otherwise a string with the BEGIN statement */
     const char* begin_statement;