From 3e61a41b9b8b44a0605773cd8a7f0dbb7d69873c Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Thu, 17 Nov 2022 17:19:37 -0800
Subject: [PATCH] Include SQL schema for CodeMirror on query pages, closes
 #1897

Refs #1893
---
 datasette/views/database.py | 25 ++++++++++++++++++++++---
 tests/test_permissions.py   | 25 +++++++++++++++++++------
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/datasette/views/database.py b/datasette/views/database.py
index 8e08c3b1..a50cf3d3 100644
--- a/datasette/views/database.py
+++ b/datasette/views/database.py
@@ -139,6 +139,9 @@ class DatabaseView(DataView):
 
         attached_databases = [d.name for d in await db.attached_databases()]
 
+        allow_execute_sql = await self.ds.permission_allowed(
+            request.actor, "execute-sql", database, default=True
+        )
         return (
             {
                 "database": database,
@@ -149,9 +152,10 @@ class DatabaseView(DataView):
                 "hidden_count": len([t for t in tables if t["hidden"]]),
                 "views": views,
                 "queries": canned_queries,
-                "allow_execute_sql": await self.ds.permission_allowed(
-                    request.actor, "execute-sql", database, default=True
-                ),
+                "allow_execute_sql": allow_execute_sql,
+                "table_columns": await _table_columns(self.ds, database)
+                if allow_execute_sql
+                else {},
             },
             {
                 "database_actions": database_actions,
@@ -508,6 +512,9 @@ class QueryView(DataView):
                 "show_hide_text": show_hide_text,
                 "show_hide_hidden": markupsafe.Markup(show_hide_hidden),
                 "hide_sql": hide_sql,
+                "table_columns": await _table_columns(self.ds, database)
+                if allow_execute_sql
+                else {},
             }
 
         return (
@@ -554,3 +561,15 @@ class MagicParameters(dict):
                     return super().__getitem__(key)
         else:
             return super().__getitem__(key)
+
+
+async def _table_columns(datasette, database_name):
+    internal = datasette.get_database("_internal")
+    result = await internal.execute(
+        "select table_name, name from columns where database_name = ?",
+        [database_name],
+    )
+    table_columns = {}
+    for row in result.rows:
+        table_columns.setdefault(row["table_name"], []).append(row["name"])
+    return table_columns
diff --git a/tests/test_permissions.py b/tests/test_permissions.py
index 8812d0f7..0364707a 100644
--- a/tests/test_permissions.py
+++ b/tests/test_permissions.py
@@ -1,7 +1,9 @@
 from .fixtures import app_client, assert_permissions_checked, make_app_client
 from bs4 import BeautifulSoup as Soup
 import copy
+import json
 import pytest
+import re
 import urllib
 
 
@@ -237,24 +239,35 @@ def test_view_query(allow, expected_anon, expected_auth):
     ],
 )
 def test_execute_sql(metadata):
+    schema_re = re.compile("const schema = ({.*?});", re.DOTALL)
     with make_app_client(metadata=metadata) as client:
         form_fragment = '<form class="sql" action="/fixtures"'
 
         # Anonymous users - should not display the form:
-        assert form_fragment not in client.get("/fixtures").text
+        anon_html = client.get("/fixtures").text
+        assert form_fragment not in anon_html
+        # And const schema should be an empty object:
+        assert "const schema = {};" in anon_html
         # This should 403:
-        assert 403 == client.get("/fixtures?sql=select+1").status
+        assert client.get("/fixtures?sql=select+1").status == 403
         # ?_where= not allowed on tables:
-        assert 403 == client.get("/fixtures/facet_cities?_where=id=3").status
+        assert client.get("/fixtures/facet_cities?_where=id=3").status == 403
 
         # But for logged in user all of these should work:
         cookies = {"ds_actor": client.actor_cookie({"id": "root"})}
         response_text = client.get("/fixtures", cookies=cookies).text
+        # Extract the schema= portion of the JavaScript
+        schema_json = schema_re.search(response_text).group(1)
+        schema = json.loads(schema_json)
+        assert set(schema["attraction_characteristic"]) == {"name", "pk"}
         assert form_fragment in response_text
-        assert 200 == client.get("/fixtures?sql=select+1", cookies=cookies).status
+        query_response = client.get("/fixtures?sql=select+1", cookies=cookies)
+        assert query_response.status == 200
+        schema2 = json.loads(schema_re.search(query_response.text).group(1))
+        assert set(schema2["attraction_characteristic"]) == {"name", "pk"}
         assert (
-            200
-            == client.get("/fixtures/facet_cities?_where=id=3", cookies=cookies).status
+            client.get("/fixtures/facet_cities?_where=id=3", cookies=cookies).status
+            == 200
         )