From b86f94883b1d827f85b07009e8346b8e9c6eeefa Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 15 Aug 2020 15:35:31 -0700
Subject: [PATCH] Don't hang in db.execute_write_fn() if connection fails

Closes #935

Refs https://github.com/simonw/latest-datasette-with-all-plugins/issues/3
---
 datasette/database.py            | 20 ++++++++++++++------
 setup.py                         |  1 +
 tests/test_internals_database.py | 17 +++++++++++++++++
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/datasette/database.py b/datasette/database.py
index ffa7a794..7ba1456b 100644
--- a/datasette/database.py
+++ b/datasette/database.py
@@ -89,14 +89,22 @@ class Database:
     def _execute_writes(self):
         # Infinite looping thread that protects the single write connection
         # to this database
-        conn = self.connect(write=True)
+        conn_exception = None
+        conn = None
+        try:
+            conn = self.connect(write=True)
+        except Exception as e:
+            conn_exception = e
         while True:
             task = self._write_queue.get()
-            try:
-                result = task.fn(conn)
-            except Exception as e:
-                print(e)
-                result = e
+            if conn_exception is not None:
+                result = conn_exception
+            else:
+                try:
+                    result = task.fn(conn)
+                except Exception as e:
+                    print(e)
+                    result = e
             task.reply_queue.sync_q.put(result)
 
     async def execute_fn(self, fn):
diff --git a/setup.py b/setup.py
index 7c352d87..5b509f8e 100644
--- a/setup.py
+++ b/setup.py
@@ -72,6 +72,7 @@ setup(
             "pytest-asyncio>=0.10,<0.15",
             "beautifulsoup4>=4.8.1,<4.10.0",
             "black~=19.10b0",
+            "pytest-timeout>=1.4.2,<1.5",
         ],
     },
     tests_require=["datasette[test]"],
diff --git a/tests/test_internals_database.py b/tests/test_internals_database.py
index 2d288cc8..832485c5 100644
--- a/tests/test_internals_database.py
+++ b/tests/test_internals_database.py
@@ -189,6 +189,23 @@ async def test_execute_write_fn_exception(db):
         await db.execute_write_fn(write_fn, block=True)
 
 
+@pytest.mark.asyncio
+@pytest.mark.timeout(1)
+async def test_execute_write_fn_connection_exception(tmpdir, app_client):
+    path = str(tmpdir / "immutable.db")
+    sqlite3.connect(path).execute("vacuum")
+    db = Database(app_client.ds, path=path, is_mutable=False)
+    app_client.ds.add_database("immutable-db", db)
+
+    def write_fn(conn):
+        assert False
+
+    with pytest.raises(AssertionError):
+        await db.execute_write_fn(write_fn, block=True)
+
+    app_client.ds.remove_database("immutable-db")
+
+
 @pytest.mark.asyncio
 async def test_mtime_ns(db):
     assert isinstance(db.mtime_ns, int)