diff --git a/tests/test_table_count.py b/tests/test_table_count.py new file mode 100644 index 00000000..135aec38 --- /dev/null +++ b/tests/test_table_count.py @@ -0,0 +1,191 @@ +import pytest +import pytest_asyncio + +from datasette import hookimpl +from datasette.app import Datasette +from datasette.database import QueryInterrupted +from datasette.filters import FilterArguments +from datasette.plugins import pm +from datasette.utils import tilde_encode + + +@pytest_asyncio.fixture +async def count_ds(tmp_path): + ds = Datasette( + [tmp_path / "counts.db"], + config={"permissions": {"execute-sql": {"id": "sql-user"}}}, + ) + db = ds.get_database() + await db.execute_write_script(""" + create table numbers (id integer primary key, name text); + insert into numbers values (1, 'one'), (2, 'two'), (3, 'three'), + (4, 'four'), (5, 'five'), (6, 'six'), (7, 'seven'); + create view number_view as select id + 0 as number from numbers; + create table "a/b.c" (id integer); + insert into "a/b.c" values (1); + create virtual table numbers_fts using fts5(name, content="numbers"); + insert into numbers_fts(numbers_fts) values ('rebuild'); + """) + db.count_limit = 2 + yield ds + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "table,query,count", + [ + ("numbers", "", 7), + ("numbers", "?id__gt=1", 6), + ("numbers", "?id__gt=1&id__lt=5", 3), + ("numbers", "?id__gt=1&id__gt=3", 4), + ("numbers", "?id__gt=100", 0), + ("numbers", "?name__contains=o", 3), + ("numbers", "?_search=two", 1), + ("numbers", "?_search_name=three", 1), + ("numbers", "?id__gt=1&_next=5&_size=1&_sort_desc=id&_nocount=1", 6), + # A computed column without affinity catches string/integer casting bugs. + ("number_view", "?number__gt=1", 6), + ("a/b.c", "", 1), + ], +) +async def test_count(count_ds, table, query, count): + response = await count_ds.client.post( + f"/counts/{tilde_encode(table)}/-/count{query}" + ) + assert response.status_code == 200, response.text + assert response.json() == {"ok": True, "count": count} + assert response.headers["cache-control"] == "no-store" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["get", "head", "put", "delete"]) +async def test_count_post_only(count_ds, method): + response = await count_ds.client.request(method, "/counts/numbers/-/count") + assert response.status_code == 405 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,status", + [ + ("/missing/numbers/-/count", 404), + ("/counts/missing/-/count", 404), + ("/counts/numbers/-/count?_fts_table=invalid", 400), + ("/counts/numbers/-/count?_where=id>1", 403), + ], +) +async def test_count_errors(count_ds, path, status): + response = await count_ds.client.post( + path, json={}, headers={"Accept": "application/json"} + ) + assert response.status_code == status, response.text + assert response.json()["ok"] is False + + +@pytest.mark.asyncio +async def test_count_permissions(tmp_path): + ds = Datasette( + [tmp_path / "counts.db"], config={"permissions": {"view-table": False}} + ) + await ds.get_database().execute_write("create table secret (id integer)") + try: + response = await ds.client.post( + "/counts/secret/-/count", json={}, headers={"Accept": "application/json"} + ) + assert response.status_code == 403 + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_count_plugin_filter(count_ds): + class Plugin: + @hookimpl + def filters_from_request(self, request): + if request.args.get("_custom"): + return FilterArguments(["id > :custom"], params={"custom": 4}) + + plugin = Plugin() + pm.register(plugin) + try: + response = await count_ds.client.post( + "/counts/numbers/-/count?_custom=1", + json={}, + headers={"Accept": "application/json"}, + ) + assert response.json() == {"ok": True, "count": 3} + finally: + pm.unregister(plugin) + + +@pytest.mark.asyncio +async def test_count_timeout(count_ds, monkeypatch): + db = count_ds.get_database() + original = db.execute + + async def execute(sql, *args, **kwargs): + if sql.startswith("select count(*)"): + raise QueryInterrupted(Exception("interrupted"), sql, {}) + return await original(sql, *args, **kwargs) + + monkeypatch.setattr(db, "execute", execute) + response = await count_ds.client.post( + "/counts/numbers/-/count", json={}, headers={"Accept": "application/json"} + ) + assert response.status_code == 400 + assert response.json()["errors"] == ["Count query timed out"] + + +@pytest.mark.asyncio +async def test_count_button_without_execute_sql(count_ds): + from bs4 import BeautifulSoup + + response = await count_ds.client.get("/counts/numbers?id__gt=1") + soup = BeautifulSoup(response.text, "html.parser") + button = soup.select_one("button.count-all") + assert button["data-count-url"] == "/counts/numbers/-/count" + assert soup.select_one(".table-count").text == "2+ rows" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query,status,count", + [ + ("?_where=id>2&_where=id<6", 200, 3), + ("?_where=invalid_sql(", 400, None), + ], +) +async def test_count_where(count_ds, query, status, count): + response = await count_ds.client.post( + "/counts/numbers/-/count" + query, + json={}, + actor={"id": "sql-user"}, + ) + assert response.status_code == status + if status == 200: + assert response.json() == {"ok": True, "count": count} + else: + assert response.json()["ok"] is False + + +@pytest.mark.asyncio +async def test_count_base_url(tmp_path): + from bs4 import BeautifulSoup + + ds = Datasette([tmp_path / "counts.db"], settings={"base_url": "/prefix/"}) + db = ds.get_database() + await db.execute_write_script( + "create table numbers (id integer); insert into numbers values (1), (2)" + ) + db.count_limit = 1 + try: + response = await ds.client.get("/counts/numbers") + button = BeautifulSoup(response.text, "html.parser").select_one(".count-all") + assert button["data-count-url"] == "/prefix/counts/numbers/-/count" + response = await ds.client.post( + ds.urls.path("/counts/numbers/-/count"), json={} + ) + assert response.json() == {"ok": True, "count": 2} + finally: + ds.close()