offset= and limit= parameters, closes #231

This commit is contained in:
Simon Willison 2021-02-14 12:02:41 -08:00
commit 320f3ac33a
4 changed files with 105 additions and 5 deletions

View file

@ -202,6 +202,12 @@ You can order all records in the table by excluding the ``where`` argument::
{'id': 1, 'age': 4, 'name': 'Cleo'}
{'id': 2, 'age': 2, 'name': 'Pancakes'}
This method also accepts ``offset=`` and ``limit=`` arguments, for specifying an OFFSET and a LIMIT for the SQL query::
>>> for row in db["dogs"].rows_where(order_by="age desc", limit=1):
... print(row)
{'id': 1, 'age': 4, 'name': 'Cleo'}
.. _python_api_get:
Retrieving a specific record
@ -1603,6 +1609,9 @@ The ``.search()`` method also accepts the following optional parameters:
``limit`` integer
Number of results to return. Defaults to all results.
``offset`` integer
Offset to use along side the limit parameter.
To return just the title and published columns for three matches for ``"dog"`` ordered by ``published`` with the most recent first, use the following:
.. code-block:: python

View file

@ -644,7 +644,15 @@ class Queryable:
def rows(self):
return self.rows_where()
def rows_where(self, where=None, where_args=None, order_by=None, select="*"):
def rows_where(
self,
where=None,
where_args=None,
order_by=None,
select="*",
limit=None,
offset=None,
):
if not self.exists():
return []
sql = "select {} from [{}]".format(select, self.name)
@ -652,6 +660,10 @@ class Queryable:
sql += " where " + where
if order_by is not None:
sql += " order by " + order_by
if limit is not None:
sql += " limit {}".format(limit)
if offset is not None:
sql += " offset {}".format(offset)
cursor = self.db.execute(sql, where_args or [])
columns = [c[0] for c in cursor.description]
for row in cursor:
@ -1454,7 +1466,7 @@ class Table(Queryable):
)
return self
def search_sql(self, columns=None, order_by=None, limit=None):
def search_sql(self, columns=None, order_by=None, limit=None, offset=None):
# Pick names for table and rank column that don't clash
original = "original_" if self.name == "original" else "original"
columns_sql = "*"
@ -1486,7 +1498,7 @@ class Table(Queryable):
[{fts_table}] match :query
order by
{order_by}
{limit}
{limit_offset}
"""
).strip()
if virtual_table_using == "FTS5":
@ -1496,6 +1508,11 @@ class Table(Queryable):
rank_implementation = "rank_bm25(matchinfo([{}], 'pcnalx'))".format(
fts_table
)
limit_offset = ""
if limit is not None:
limit_offset += " limit {}".format(limit)
if offset is not None:
limit_offset += " offset {}".format(offset)
return sql.format(
dbtable=self.name,
original=original,
@ -1503,15 +1520,16 @@ class Table(Queryable):
columns_with_prefix=columns_with_prefix_sql,
fts_table=fts_table,
order_by=order_by or rank_implementation,
limit="limit {}".format(limit) if limit else "",
limit_offset=limit_offset.strip(),
).strip()
def search(self, q, order_by=None, columns=None, limit=None):
def search(self, q, order_by=None, columns=None, limit=None, offset=None):
cursor = self.db.execute(
self.search_sql(
order_by=order_by,
columns=columns,
limit=limit,
offset=offset,
),
{"query": q},
)

View file

@ -82,6 +82,18 @@ def test_enable_fts_escape_table_names(fresh_db):
assert [] == list(table.search("bar"))
def test_search_limit_offset(fresh_db):
table = fresh_db["t"]
table.insert_all(search_records)
table.enable_fts(["text", "country"], fts_version="FTS4")
assert len(list(table.search("are"))) == 2
assert len(list(table.search("are", limit=1))) == 1
assert list(table.search("are", limit=1, order_by="rowid"))[0]["rowid"] == 1
assert (
list(table.search("are", limit=1, offset=1, order_by="rowid"))[0]["rowid"] == 2
)
def test_enable_fts_table_names_containing_spaces(fresh_db):
table = fresh_db["test"]
table.insert({"column with spaces": "in its name"})
@ -424,6 +436,50 @@ def test_enable_fts_replace_does_nothing_if_args_the_same():
" rank_bm25(matchinfo([books_fts], 'pcnalx'))"
),
),
(
{"offset": 1, "limit": 1},
"FTS4",
(
"with original as (\n"
" select\n"
" rowid,\n"
" *\n"
" from [books]\n"
")\n"
"select\n"
" [original].*\n"
"from\n"
" [original]\n"
" join [books_fts] on [original].rowid = [books_fts].rowid\n"
"where\n"
" [books_fts] match :query\n"
"order by\n"
" rank_bm25(matchinfo([books_fts], 'pcnalx'))\n"
"limit 1 offset 1"
),
),
(
{"limit": 2},
"FTS4",
(
"with original as (\n"
" select\n"
" rowid,\n"
" *\n"
" from [books]\n"
")\n"
"select\n"
" [original].*\n"
"from\n"
" [original]\n"
" join [books_fts] on [original].rowid = [books_fts].rowid\n"
"where\n"
" [books_fts] match :query\n"
"order by\n"
" rank_bm25(matchinfo([books_fts], 'pcnalx'))\n"
"limit 2"
),
),
],
)
def test_search_sql(kwargs, fts, expected):

View file

@ -51,3 +51,20 @@ def test_rows_where_order_by(where, order_by, expected_ids, fresh_db):
pk="id",
)
assert expected_ids == [r["id"] for r in table.rows_where(where, order_by=order_by)]
@pytest.mark.parametrize(
"offset,limit,expected",
[
(None, 3, [1, 2, 3]),
(0, 3, [1, 2, 3]),
(3, 3, [4, 5, 6]),
],
)
def test_rows_where_offset_limit(fresh_db, offset, limit, expected):
table = fresh_db["rows"]
table.insert_all([{"id": id} for id in range(1, 101)], pk="id")
assert table.count == 100
assert expected == [
r["id"] for r in table.rows_where(offset=offset, limit=limit, order_by="id")
]