From 9cb69cbd45ed8fd93190c47060c19abec80bc4ef Mon Sep 17 00:00:00 2001
From: Simon Willison <simonw@eventbrite.com>
Date: Tue, 14 Nov 2017 18:55:10 -0800
Subject: [PATCH] New ?_sql_time_limit_ms=10 argument to database and table
 page

Allows callers to opt for a lower time limit.

Closes #95
---
 datasette/app.py  | 23 +++++++++++++++++++----
 tests/test_app.py | 12 ++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index 853f24a0..ff3de281 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -114,7 +114,7 @@ class BaseView(HTTPMethodView):
         for name, num_args, func in self.ds.sqlite_functions:
             conn.create_function(name, num_args, func)
 
-    async def execute(self, db_name, sql, params=None, truncate=False):
+    async def execute(self, db_name, sql, params=None, truncate=False, custom_time_limit=None):
         """Executes sql against db_name in a thread"""
         def sql_operation_in_thread():
             conn = getattr(connections, db_name, None)
@@ -128,7 +128,11 @@ class BaseView(HTTPMethodView):
                 self.prepare_connection(conn)
                 setattr(connections, db_name, conn)
 
-            with sqlite_timelimit(conn, self.ds.sql_time_limit_ms):
+            time_limit_ms = self.ds.sql_time_limit_ms
+            if custom_time_limit and custom_time_limit < self.ds.sql_time_limit_ms:
+                time_limit_ms = custom_time_limit
+
+            with sqlite_timelimit(conn, time_limit_ms):
                 try:
                     cursor = conn.cursor()
                     cursor.execute(sql, params or {})
@@ -312,7 +316,12 @@ class DatabaseView(BaseView):
         params = request.raw_args
         sql = params.pop('sql')
         validate_sql_select(sql)
-        rows, truncated, description = await self.execute(name, sql, params, truncate=True)
+        extra_args = {}
+        if params.get('_sql_time_limit_ms'):
+            extra_args['custom_time_limit'] = int(params['_sql_time_limit_ms'])
+        rows, truncated, description = await self.execute(
+            name, sql, params, truncate=True, **extra_args
+        )
         columns = [r[0] for r in description]
         return {
             'database': name,
@@ -428,7 +437,13 @@ class TableView(BaseView):
             offset=offset,
         )
 
-        rows, truncated, description = await self.execute(name, sql, params, truncate=True)
+        extra_args = {}
+        if request.raw_args.get('_sql_time_limit_ms'):
+            extra_args['custom_time_limit'] = int(request.raw_args['_sql_time_limit_ms'])
+
+        rows, truncated, description = await self.execute(
+            name, sql, params, truncate=True, **extra_args
+        )
 
         columns = [r[0] for r in description]
         display_columns = columns
diff --git a/tests/test_app.py b/tests/test_app.py
index ff910c3c..f25a9f0d 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -101,6 +101,18 @@ def test_sql_time_limit(app_client):
     assert 'interrupted' == response.json['error']
 
 
+def test_custom_sql_time_limit(app_client):
+    _, response = app_client.get(
+        '/test_tables.jsono?sql=select+sleep(0.01)'
+    )
+    assert 200 == response.status
+    _, response = app_client.get(
+        '/test_tables.jsono?sql=select+sleep(0.01)&_sql_time_limit_ms=5'
+    )
+    assert 400 == response.status
+    assert 'interrupted' == response.json['error']
+
+
 def test_invalid_custom_sql(app_client):
     _, response = app_client.get(
         '/test_tables?sql=.schema'