From 7ecd3425d45a37efbc745788dfe21df0286c785a Mon Sep 17 00:00:00 2001
From: Erlend Egeberg Aasland <erlend.aasland@innova.no>
Date: Wed, 25 Aug 2021 12:59:42 +0200
Subject: [PATCH] bpo-27334: roll back transaction if sqlite3 context manager
 fails to commit (GH-26202)

Co-authored-by: Luca Citi
Co-authored-by: Berker Peksag <berker.peksag@gmail.com>
---
 Lib/sqlite3/test/dbapi.py                     | 80 ++++++++++++++++++-
 .../2021-05-18-00-17-21.bpo-27334.32EJZi.rst  |  2 +
 Modules/_sqlite/connection.c                  | 26 ++++--
 3 files changed, 101 insertions(+), 7 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2021-05-18-00-17-21.bpo-27334.32EJZi.rst

diff --git a/Lib/sqlite3/test/dbapi.py b/Lib/sqlite3/test/dbapi.py
index 5d7e5bba05b..bb9d5a7ce3e 100644
--- a/Lib/sqlite3/test/dbapi.py
+++ b/Lib/sqlite3/test/dbapi.py
@@ -22,11 +22,17 @@
 
 import contextlib
 import sqlite3 as sqlite
+import subprocess
 import sys
 import threading
 import unittest
 
-from test.support import check_disallow_instantiation, threading_helper, bigmemtest
+from test.support import (
+    bigmemtest,
+    check_disallow_instantiation,
+    threading_helper,
+    SHORT_TIMEOUT,
+)
 from test.support.os_helper import TESTFN, unlink
 
 
@@ -986,6 +992,77 @@ def test_on_conflict_replace(self):
         self.assertEqual(self.cu.fetchall(), [('Very different data!', 'foo')])
 
 
+class MultiprocessTests(unittest.TestCase):
+    CONNECTION_TIMEOUT = SHORT_TIMEOUT / 1000.  # Defaults to 30 ms
+
+    def tearDown(self):
+        unlink(TESTFN)
+
+    def test_ctx_mgr_rollback_if_commit_failed(self):
+        # bpo-27334: ctx manager does not rollback if commit fails
+        SCRIPT = f"""if 1:
+            import sqlite3
+            def wait():
+                print("started")
+                assert "database is locked" in input()
+
+            cx = sqlite3.connect("{TESTFN}", timeout={self.CONNECTION_TIMEOUT})
+            cx.create_function("wait", 0, wait)
+            with cx:
+                cx.execute("create table t(t)")
+            try:
+                # execute two transactions; both will try to lock the db
+                cx.executescript('''
+                    -- start a transaction and wait for parent
+                    begin transaction;
+                    select * from t;
+                    select wait();
+                    rollback;
+
+                    -- start a new transaction; would fail if parent holds lock
+                    begin transaction;
+                    select * from t;
+                    rollback;
+                ''')
+            finally:
+                cx.close()
+        """
+
+        # spawn child process
+        proc = subprocess.Popen(
+            [sys.executable, "-c", SCRIPT],
+            encoding="utf-8",
+            bufsize=0,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+        )
+        self.addCleanup(proc.communicate)
+
+        # wait for child process to start
+        self.assertEqual("started", proc.stdout.readline().strip())
+
+        cx = sqlite.connect(TESTFN, timeout=self.CONNECTION_TIMEOUT)
+        try:  # context manager should correctly release the db lock
+            with cx:
+                cx.execute("insert into t values('test')")
+        except sqlite.OperationalError as exc:
+            proc.stdin.write(str(exc))
+        else:
+            proc.stdin.write("no error")
+        finally:
+            cx.close()
+
+        # terminate child process
+        self.assertIsNone(proc.returncode)
+        try:
+            proc.communicate(input="end", timeout=SHORT_TIMEOUT)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            proc.communicate()
+            raise
+        self.assertEqual(proc.returncode, 0)
+
+
 def suite():
     tests = [
         ClosedConTests,
@@ -995,6 +1072,7 @@ def suite():
         CursorTests,
         ExtensionTests,
         ModuleTests,
+        MultiprocessTests,
         OpenTests,
         SqliteOnConflictTests,
         ThreadTests,
diff --git a/Misc/NEWS.d/next/Library/2021-05-18-00-17-21.bpo-27334.32EJZi.rst b/Misc/NEWS.d/next/Library/2021-05-18-00-17-21.bpo-27334.32EJZi.rst
new file mode 100644
index 00000000000..dc0cdf33ec5
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-05-18-00-17-21.bpo-27334.32EJZi.rst
@@ -0,0 +1,2 @@
+The :mod:`sqlite3` context manager now performs a rollback (thus releasing the
+database lock) if commit failed.  Patch by Luca Citi and Erlend E. Aasland.
diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c
index 8ad5f5f061d..1bc045523a2 100644
--- a/Modules/_sqlite/connection.c
+++ b/Modules/_sqlite/connection.c
@@ -1785,17 +1785,31 @@ pysqlite_connection_exit_impl(pysqlite_Connection *self, PyObject *exc_type,
                               PyObject *exc_value, PyObject *exc_tb)
 /*[clinic end generated code: output=0705200e9321202a input=bd66f1532c9c54a7]*/
 {
-    const char* method_name;
+    int commit = 0;
     PyObject* result;
 
     if (exc_type == Py_None && exc_value == Py_None && exc_tb == Py_None) {
-        method_name = "commit";
-    } else {
-        method_name = "rollback";
+        commit = 1;
+        result = pysqlite_connection_commit_impl(self);
+    }
+    else {
+        result = pysqlite_connection_rollback_impl(self);
     }
 
-    result = PyObject_CallMethod((PyObject*)self, method_name, NULL);
-    if (!result) {
+    if (result == NULL) {
+        if (commit) {
+            /* Commit failed; try to rollback in order to unlock the database.
+             * If rollback also fails, chain the exceptions. */
+            PyObject *exc, *val, *tb;
+            PyErr_Fetch(&exc, &val, &tb);
+            result = pysqlite_connection_rollback_impl(self);
+            if (result == NULL) {
+                _PyErr_ChainExceptions(exc, val, tb);
+            }
+            else {
+                PyErr_Restore(exc, val, tb);
+            }
+        }
         return NULL;
     }
     Py_DECREF(result);