datasette/tests/test_table_count.py

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()