diff --git a/aiosqlite/core.py b/aiosqlite/core.py index 3deaa77..60db6a1 100644 --- a/aiosqlite/core.py +++ b/aiosqlite/core.py @@ -129,6 +129,19 @@ def _conn(self) -> sqlite3.Connection: return self._connection + def _execute_insert( + self, sql: str, parameters: Iterable[Any] + ) -> Optional[sqlite3.Row]: + cursor = self._conn.execute(sql, parameters) + cursor.execute("SELECT last_insert_rowid()") + return cursor.fetchone() + + def _execute_fetchall( + self, sql: str, parameters: Iterable[Any] + ) -> Iterable[sqlite3.Row]: + cursor = self._conn.execute(sql, parameters) + return cursor.fetchall() + def run(self) -> None: """Execute function calls on a separate thread.""" while self._running: @@ -211,6 +224,24 @@ async def execute(self, sql: str, parameters: Iterable[Any] = None) -> Cursor: cursor = await self._execute(self._conn.execute, sql, parameters) return Cursor(self, cursor) + @contextmanager + async def execute_insert( + self, sql: str, parameters: Iterable[Any] = None + ) -> Optional[sqlite3.Row]: + """Helper to insert and get the last_insert_rowid.""" + if parameters is None: + parameters = [] + return await self._execute(self._execute_insert, sql, parameters) + + @contextmanager + async def execute_fetchall( + self, sql: str, parameters: Iterable[Any] = None + ) -> Iterable[sqlite3.Row]: + """Helper to execute a query and return all the data.""" + if parameters is None: + parameters = [] + return await self._execute(self._execute_fetchall, sql, parameters) + @contextmanager async def executemany( self, sql: str, parameters: Iterable[Iterable[Any]] diff --git a/tests/perf.py b/tests/perf.py index e357435..c1c830c 100644 --- a/tests/perf.py +++ b/tests/perf.py @@ -103,3 +103,52 @@ async def test_inserts(self): yield await db.execute("insert into perf (k) values (1), (2), (3)") await db.commit() + + @timed + async def test_insert_ids(self): + async with aiosqlite.connect(TEST_DB) as db: + await db.execute("create table perf (i integer primary key asc, k integer)") + await db.commit() + + while True: + yield + cursor = await db.execute("insert into perf (k) values (1)") + await cursor.execute("select last_insert_rowid()") + await cursor.fetchone() + await db.commit() + + @timed + async def test_insert_macro_ids(self): + async with aiosqlite.connect(TEST_DB) as db: + await db.execute("create table perf (i integer primary key asc, k integer)") + await db.commit() + + while True: + yield + await db.execute_insert("insert into perf (k) values (1)") + await db.commit() + + @timed + async def test_select(self): + async with aiosqlite.connect(TEST_DB) as db: + await db.execute("create table perf (i integer primary key asc, k integer)") + for i in range(100): + await db.execute("insert into perf (k) values (%d)" % (i, )) + await db.commit() + + while True: + yield + cursor = await db.execute("select i, k from perf") + assert len(await cursor.fetchall()) == 100 + + @timed + async def test_select_macro(self): + async with aiosqlite.connect(TEST_DB) as db: + await db.execute("create table perf (i integer primary key asc, k integer)") + for i in range(100): + await db.execute("insert into perf (k) values (%d)" % (i, )) + await db.commit() + + while True: + yield + assert len(await db.execute_fetchall("select i, k from perf")) == 100