From 320f3ac33a83b32f89559ef0c162b7eca428a278 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 14 Feb 2021 12:02:41 -0800 Subject: [PATCH] offset= and limit= parameters, closes #231 --- docs/python-api.rst | 9 ++++++++ sqlite_utils/db.py | 28 +++++++++++++++++++---- tests/test_fts.py | 56 +++++++++++++++++++++++++++++++++++++++++++++ tests/test_rows.py | 17 ++++++++++++++ 4 files changed, 105 insertions(+), 5 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index c656045..246a301 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 14a9efc..17d79de 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -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}, ) diff --git a/tests/test_fts.py b/tests/test_fts.py index 06fac20..94650bf 100644 --- a/tests/test_fts.py +++ b/tests/test_fts.py @@ -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): diff --git a/tests/test_rows.py b/tests/test_rows.py index 603f254..6d995e2 100644 --- a/tests/test_rows.py +++ b/tests/test_rows.py @@ -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") + ]