From 55a709c480a1e7401b4ff6208f37a2cf7c682183 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 26 Oct 2022 14:34:33 -0700
Subject: [PATCH] Allow leading comments on SQL queries, refs #1860

---
 datasette/utils/__init__.py | 27 +++++++++++++++++++++------
 tests/test_utils.py         |  7 +++++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/datasette/utils/__init__.py b/datasette/utils/__init__.py
index 803ba96d..977a66d6 100644
--- a/datasette/utils/__init__.py
+++ b/datasette/utils/__init__.py
@@ -205,13 +205,28 @@ class InvalidSql(Exception):
     pass
 
 
+# Allow SQL to start with a /* */ or -- comment
+comment_re = (
+    # Start of string, then any amount of whitespace
+    r"^(\s*"
+    +
+    # Comment that starts with -- and ends at a newline
+    r"(?:\-\-.*?\n\s*)"
+    +
+    # Comment that starts with /* and ends with */
+    r"|(?:/\*[\s\S]*?\*/)"
+    +
+    # Whitespace
+    r")*\s*"
+)
+
 allowed_sql_res = [
-    re.compile(r"^select\b"),
-    re.compile(r"^explain\s+select\b"),
-    re.compile(r"^explain\s+query\s+plan\s+select\b"),
-    re.compile(r"^with\b"),
-    re.compile(r"^explain\s+with\b"),
-    re.compile(r"^explain\s+query\s+plan\s+with\b"),
+    re.compile(comment_re + r"select\b"),
+    re.compile(comment_re + r"explain\s+select\b"),
+    re.compile(comment_re + r"explain\s+query\s+plan\s+select\b"),
+    re.compile(comment_re + r"with\b"),
+    re.compile(comment_re + r"explain\s+with\b"),
+    re.compile(comment_re + r"explain\s+query\s+plan\s+with\b"),
 ]
 allowed_pragmas = (
     "database_list",
diff --git a/tests/test_utils.py b/tests/test_utils.py
index d71a612d..e89f1e6b 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -141,6 +141,7 @@ def test_custom_json_encoder(obj, expected):
         "update blah set some_column='# Hello there\n\n* This is a list\n* of items\n--\n[And a link](https://github.com/simonw/datasette-render-markdown).'\nas demo_markdown",
         "PRAGMA case_sensitive_like = true",
         "SELECT * FROM pragma_not_on_allow_list('idx52')",
+        "/* This comment is not valid. select 1",
     ],
 )
 def test_validate_sql_select_bad(bad_sql):
@@ -166,6 +167,12 @@ def test_validate_sql_select_bad(bad_sql):
         "explain query plan WITH RECURSIVE cnt(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM cnt LIMIT 10) SELECT x FROM cnt;",
         "SELECT * FROM pragma_index_info('idx52')",
         "select * from pragma_table_xinfo('table')",
+        # Various types of comment
+        "-- comment\nselect 1",
+        "-- one line\n  -- two line\nselect 1",
+        "  /* comment */\nselect 1",
+        "  /* comment */select 1",
+        "/* comment */\n -- another\n /* one more */ select 1",
     ],
 )
 def test_validate_sql_select_good(good_sql):