From 5cadc244895fc47e0534c6e90df976d34293921e Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sat, 18 Dec 2021 10:57:22 -0800 Subject: [PATCH] db.execute_write_script() and db.execute_write_many(), closes #1570 Refs #1555 --- datasette/database.py | 29 ++++++++++++++++++++--------- datasette/utils/internal_db.py | 2 +- docs/internals.rst | 26 +++++++++++++++++++++++--- tests/test_internals_database.py | 24 ++++++++++++++---------- 4 files changed, 58 insertions(+), 23 deletions(-) diff --git a/datasette/database.py b/datasette/database.py index f8365f5c08..1de1d5ec10 100644 --- a/datasette/database.py +++ b/datasette/database.py @@ -94,22 +94,33 @@ def connect(self, write=False): f"file:{self.path}{qs}", uri=True, check_same_thread=False ) - async def execute_write(self, sql, params=None, executescript=False, block=False): - assert not ( - executescript and params - ), "Cannot use params with executescript=True" - + async def execute_write(self, sql, params=None, block=False): def _inner(conn): with conn: - if executescript: - return conn.executescript(sql) - else: - return conn.execute(sql, params or []) + return conn.execute(sql, params or []) with trace("sql", database=self.name, sql=sql.strip(), params=params): results = await self.execute_write_fn(_inner, block=block) return results + async def execute_write_script(self, sql, block=False): + def _inner(conn): + with conn: + return conn.executescript(sql) + + with trace("sql", database=self.name, sql=sql.strip(), executescript=True): + results = await self.execute_write_fn(_inner, block=block) + return results + + async def execute_write_many(self, sql, params_seq, block=False): + def _inner(conn): + with conn: + return conn.executemany(sql, params_seq) + + with trace("sql", database=self.name, sql=sql.strip(), executemany=True): + results = await self.execute_write_fn(_inner, block=block) + return results + async def execute_write_fn(self, fn, block=False): task_id = uuid.uuid5(uuid.NAMESPACE_DNS, "datasette.io") if self._write_queue is None: diff --git a/datasette/utils/internal_db.py b/datasette/utils/internal_db.py index 80babff85a..8a145767f2 100644 --- a/datasette/utils/internal_db.py +++ b/datasette/utils/internal_db.py @@ -61,7 +61,7 @@ async def init_internal_db(db): ); """ ).strip() - await db.execute_write(create_tables_sql, block=True, executescript=True) + await db.execute_write_script(create_tables_sql, block=True) async def populate_schema_tables(internal_db, db): diff --git a/docs/internals.rst b/docs/internals.rst index d40e679b26..bc0174a822 100644 --- a/docs/internals.rst +++ b/docs/internals.rst @@ -663,8 +663,8 @@ Example usage: .. _database_execute_write: -await db.execute_write(sql, params=None, executescript=False, block=False) --------------------------------------------------------------------------- +await db.execute_write(sql, params=None, block=False) +----------------------------------------------------- SQLite only allows one database connection to write at a time. Datasette handles this for you by maintaining a queue of writes to be executed against a given database. Plugins can submit write operations to this queue and they will be executed in the order in which they are received. @@ -676,7 +676,27 @@ By default queries are considered to be "fire and forget" - they will be added t If you pass ``block=True`` this behaviour changes: the method will block until the write operation has completed, and the return value will be the return from calling ``conn.execute(...)`` using the underlying ``sqlite3`` Python library. -If you pass ``executescript=True`` your SQL will be executed using the ``sqlite3`` `conn.executescript() `__ method. This allows multiple SQL statements to be separated by semicolons, but cannot be used with the ``params=`` option. +.. _database_execute_write_script: + +await db.execute_write_script(sql, block=False) +----------------------------------------------- + +Like ``execute_write()`` but can be used to send multiple SQL statements in a single string separated by semicolons, using the ``sqlite3`` `conn.executescript() `__ method. + +.. _database_execute_write_many: + +await db.execute_write_many(sql, params_seq, block=False) +--------------------------------------------------------- + +Like ``execute_write()`` but uses the ``sqlite3`` `conn.executemany() `__ method. This will efficiently execute the same SQL statement against each of the parameters in the ``params_seq`` iterator, for example: + +.. code-block:: python + + await db.execute_write_many( + "insert into characters (id, name) values (?, ?)", + [(1, "Melanie"), (2, "Selma"), (2, "Viktor")], + block=True, + ) .. _database_execute_write_fn: diff --git a/tests/test_internals_database.py b/tests/test_internals_database.py index aa5676e769..f751bf9ac1 100644 --- a/tests/test_internals_database.py +++ b/tests/test_internals_database.py @@ -397,10 +397,9 @@ async def test_execute_write_block_false(db): @pytest.mark.asyncio -async def test_execute_write_executescript(db): - await db.execute_write( +async def test_execute_write_script(db): + await db.execute_write_script( "create table foo (id integer primary key); create table bar (id integer primary key); ", - executescript=True, block=True, ) table_names = await db.table_names() @@ -408,13 +407,18 @@ async def test_execute_write_executescript(db): @pytest.mark.asyncio -async def test_execute_write_executescript_not_allowed_with_params(db): - with pytest.raises(AssertionError): - await db.execute_write( - "update roadside_attractions set name = ? where pk = ?", - ["Mystery!", 1], - executescript=True, - ) +async def test_execute_write_many(db): + await db.execute_write_script( + "create table foomany (id integer primary key)", + block=True, + ) + await db.execute_write_many( + "insert into foomany (id) values (?)", + [(1,), (10,), (100,)], + block=True, + ) + result = await db.execute("select * from foomany") + assert [r[0] for r in result.rows] == [1, 10, 100] @pytest.mark.asyncio