From 7e1abd0da4cc57be64685b6ea717457ceefbabb3 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 24 May 2026 22:37:34 -0700 Subject: [PATCH] Add internal query storage APIs Refs #2735 --- datasette/app.py | 200 +++++++++++++++++++++++++++++++++ datasette/utils/internal_db.py | 28 +++++ tests/test_queries.py | 123 ++++++++++++++++++++ 3 files changed, 351 insertions(+) create mode 100644 tests/test_queries.py diff --git a/datasette/app.py b/datasette/app.py index 75f05d88..518215fd 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -268,6 +268,7 @@ DEFAULT_SETTINGS = {option.name: option.default for option in SETTINGS} FAVICON_PATH = app_root / "datasette" / "static" / "favicon.png" DEFAULT_NOT_SET = object() +UNCHANGED = object() ResourcesSQL = collections.namedtuple("ResourcesSQL", ("sql", "params")) @@ -1007,6 +1008,205 @@ class Datasette: [database_name, resource_name, column_name, key, value], ) + @staticmethod + def _query_row_to_dict(row): + if row is None: + return None + parameters = json.loads(row["parameters"] or "[]") + is_write = bool(row["is_write"]) + return { + "database": row["database_name"], + "name": row["name"], + "sql": row["sql"], + "title": row["title"], + "description": row["description"], + "description_html": row["description_html"], + "hide_sql": bool(row["hide_sql"]), + "fragment": row["fragment"], + "params": parameters, + "parameters": parameters, + "is_write": is_write, + "write": is_write, + "published": bool(row["published"]), + "source": row["source"], + "owner_id": row["owner_id"], + "on_success_message": row["on_success_message"], + "on_success_message_sql": row["on_success_message_sql"], + "on_success_redirect": row["on_success_redirect"], + "on_error_message": row["on_error_message"], + "on_error_redirect": row["on_error_redirect"], + } + + async def add_query( + self, + database, + name, + sql, + *, + title=None, + description=None, + description_html=None, + hide_sql=False, + fragment=None, + parameters=None, + is_write=False, + published=False, + source="plugin", + owner_id=None, + on_success_message=None, + on_success_message_sql=None, + on_success_redirect=None, + on_error_message=None, + on_error_redirect=None, + replace=True, + ): + parameters_json = json.dumps(list(parameters or [])) + sql_statement = """ + INSERT INTO queries ( + database_name, name, sql, title, description, description_html, + hide_sql, fragment, parameters, is_write, published, source, + owner_id, on_success_message, on_success_message_sql, + on_success_redirect, on_error_message, on_error_redirect + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + if replace: + sql_statement += """ + ON CONFLICT(database_name, name) DO UPDATE SET + sql = excluded.sql, + title = excluded.title, + description = excluded.description, + description_html = excluded.description_html, + hide_sql = excluded.hide_sql, + fragment = excluded.fragment, + parameters = excluded.parameters, + is_write = excluded.is_write, + published = excluded.published, + source = excluded.source, + owner_id = excluded.owner_id, + on_success_message = excluded.on_success_message, + on_success_message_sql = excluded.on_success_message_sql, + on_success_redirect = excluded.on_success_redirect, + on_error_message = excluded.on_error_message, + on_error_redirect = excluded.on_error_redirect, + updated_at = CURRENT_TIMESTAMP + """ + await self.get_internal_database().execute_write( + sql_statement, + [ + database, + name, + sql, + title, + description, + description_html, + int(bool(hide_sql)), + fragment, + parameters_json, + int(bool(is_write)), + int(bool(published)), + source, + owner_id, + on_success_message, + on_success_message_sql, + on_success_redirect, + on_error_message, + on_error_redirect, + ], + ) + + async def update_query( + self, + database, + name, + *, + sql=UNCHANGED, + title=UNCHANGED, + description=UNCHANGED, + description_html=UNCHANGED, + hide_sql=UNCHANGED, + fragment=UNCHANGED, + parameters=UNCHANGED, + is_write=UNCHANGED, + published=UNCHANGED, + source=UNCHANGED, + owner_id=UNCHANGED, + on_success_message=UNCHANGED, + on_success_message_sql=UNCHANGED, + on_success_redirect=UNCHANGED, + on_error_message=UNCHANGED, + on_error_redirect=UNCHANGED, + ): + fields = { + "sql": sql, + "title": title, + "description": description, + "description_html": description_html, + "hide_sql": hide_sql, + "fragment": fragment, + "parameters": parameters, + "is_write": is_write, + "published": published, + "source": source, + "owner_id": owner_id, + "on_success_message": on_success_message, + "on_success_message_sql": on_success_message_sql, + "on_success_redirect": on_success_redirect, + "on_error_message": on_error_message, + "on_error_redirect": on_error_redirect, + } + updates = [] + params = [] + for field, value in fields.items(): + if value is UNCHANGED: + continue + if field in {"hide_sql", "is_write", "published"}: + value = int(bool(value)) + elif field == "parameters": + value = json.dumps(list(value or [])) + updates.append(f"{field} = ?") + params.append(value) + if not updates: + return + updates.append("updated_at = CURRENT_TIMESTAMP") + params.extend([database, name]) + await self.get_internal_database().execute_write( + """ + UPDATE queries + SET {} + WHERE database_name = ? AND name = ? + """.format(", ".join(updates)), + params, + ) + + async def remove_query(self, database, name, source=None): + sql = "DELETE FROM queries WHERE database_name = ? AND name = ?" + params = [database, name] + if source is not None: + sql += " AND source = ?" + params.append(source) + await self.get_internal_database().execute_write(sql, params) + + async def get_query(self, database, name): + rows = await self.get_internal_database().execute( + """ + SELECT * FROM queries + WHERE database_name = ? AND name = ? + """, + [database, name], + ) + return self._query_row_to_dict(rows.first()) + + async def get_queries(self, database): + rows = await self.get_internal_database().execute( + """ + SELECT * FROM queries + WHERE database_name = ? + ORDER BY name + """, + [database], + ) + return {row["name"]: self._query_row_to_dict(row) for row in rows} + # Column types API async def _get_resource_column_details(self, database: str, resource: str): diff --git a/datasette/utils/internal_db.py b/datasette/utils/internal_db.py index df149928..9008c083 100644 --- a/datasette/utils/internal_db.py +++ b/datasette/utils/internal_db.py @@ -112,6 +112,34 @@ async def initialize_metadata_tables(db): config TEXT, PRIMARY KEY (database_name, resource_name, column_name) ); + + CREATE TABLE IF NOT EXISTS queries ( + database_name TEXT NOT NULL, + name TEXT NOT NULL, + sql TEXT NOT NULL, + title TEXT, + description TEXT, + description_html TEXT, + hide_sql INTEGER NOT NULL DEFAULT 0 CHECK (hide_sql IN (0, 1)), + fragment TEXT, + parameters TEXT NOT NULL DEFAULT '[]', + is_write INTEGER NOT NULL DEFAULT 0 CHECK (is_write IN (0, 1)), + published INTEGER NOT NULL DEFAULT 0 CHECK (published IN (0, 1)), + source TEXT NOT NULL DEFAULT 'user', + owner_id TEXT, + on_success_message TEXT, + on_success_message_sql TEXT, + on_success_redirect TEXT, + on_error_message TEXT, + on_error_redirect TEXT, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (database_name, name), + CHECK (is_write = 0 OR published = 0) + ); + + CREATE INDEX IF NOT EXISTS queries_owner_idx + ON queries(owner_id); """)) diff --git a/tests/test_queries.py b/tests/test_queries.py new file mode 100644 index 00000000..d30fcfe7 --- /dev/null +++ b/tests/test_queries.py @@ -0,0 +1,123 @@ +import pytest + +from datasette.app import Datasette + + +@pytest.mark.asyncio +async def test_queries_internal_table_schema(): + ds = Datasette(memory=True) + await ds.invoke_startup() + internal_db = ds.get_internal_database() + + columns = [ + row["name"] + for row in ( + await internal_db.execute("select name from pragma_table_info('queries')") + ) + ] + + assert columns == [ + "database_name", + "name", + "sql", + "title", + "description", + "description_html", + "hide_sql", + "fragment", + "parameters", + "is_write", + "published", + "source", + "owner_id", + "on_success_message", + "on_success_message_sql", + "on_success_redirect", + "on_error_message", + "on_error_redirect", + "created_at", + "updated_at", + ] + + +@pytest.mark.asyncio +async def test_add_get_and_remove_query(): + ds = Datasette(memory=True) + ds.add_memory_database("query_api", name="data") + await ds.invoke_startup() + + await ds.add_query( + "data", + "top_customers", + "select * from customers where region = :region", + title="Top customers", + description="Customers by region", + hide_sql=True, + fragment="chart", + parameters=["region"], + published=True, + source="user", + owner_id="alice", + ) + + query = await ds.get_query("data", "top_customers") + assert query == { + "database": "data", + "name": "top_customers", + "sql": "select * from customers where region = :region", + "title": "Top customers", + "description": "Customers by region", + "description_html": None, + "hide_sql": True, + "fragment": "chart", + "params": ["region"], + "parameters": ["region"], + "is_write": False, + "write": False, + "published": True, + "source": "user", + "owner_id": "alice", + "on_success_message": None, + "on_success_message_sql": None, + "on_success_redirect": None, + "on_error_message": None, + "on_error_redirect": None, + } + + assert await ds.get_queries("data") == {"top_customers": query} + + await ds.remove_query("data", "top_customers") + assert await ds.get_query("data", "top_customers") is None + assert await ds.get_queries("data") == {} + + +@pytest.mark.asyncio +async def test_update_query_only_updates_provided_fields(): + ds = Datasette(memory=True) + ds.add_memory_database("query_api_update", name="data") + await ds.invoke_startup() + + await ds.add_query( + "data", + "redirect", + "select 1", + title="Original", + on_success_redirect="/original", + parameters=["one"], + ) + + await ds.update_query( + "data", + "redirect", + title="Updated", + parameters=[], + on_success_redirect=None, + ) + + query = await ds.get_query("data", "redirect") + assert query["title"] == "Updated" + assert query["parameters"] == [] + assert query["params"] == [] + assert query["on_success_redirect"] is None + assert query["sql"] == "select 1" + assert query["published"] is False