Merge pull request #43 from encode/return-val-from-execute

Return result from execute

Author: lovelydinosaur

databases/__init__.py

@@ -1,5 +1,5 @@
 from databases.core import Database, DatabaseURL
 
 
-__version__ = "0.1.4"
+__version__ = "0.1.5"
 __all__ = ["Database", "DatabaseURL"]

databases/backends/mysql.py

@@ -92,14 +92,15 @@ async def fetch_one(self, query: ClauseElement) -> RowProxy:
         finally:
             await cursor.close()
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> None:
+    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
         assert self._connection is not None, "Connection is not acquired"
         if values is not None:
             query = query.values(values)
         query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
             await cursor.execute(query, args)
+            return cursor.lastrowid
         finally:
             await cursor.close()
 

databases/backends/postgres.py

@@ -99,12 +99,12 @@ async def fetch_one(self, query: ClauseElement) -> typing.Any:
         row = await self._connection.fetchrow(query, *args)
         return Record(row, result_columns, self._dialect)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> None:
+    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
         assert self._connection is not None, "Connection is not acquired"
         if values is not None:
             query = query.values(values)
         query, args, result_columns = self._compile(query)
-        await self._connection.execute(query, *args)
+        return await self._connection.fetchval(query, *args)
 
     async def execute_many(self, query: ClauseElement, values: list) -> None:
         assert self._connection is not None, "Connection is not acquired"

databases/backends/sqlite.py

@@ -107,6 +107,7 @@ async def execute(self, query: ClauseElement, values: dict = None) -> None:
         query, args, context = self._compile(query)
         cursor = await self._connection.execute(query, args)
         await cursor.close()
+        return cursor.lastrowid
 
     async def execute_many(self, query: ClauseElement, values: list) -> None:
         assert self._connection is not None, "Connection is not acquired"

databases/core.py

@@ -96,7 +96,7 @@ async def fetch_one(self, query: ClauseElement) -> RowProxy:
         async with self.connection() as connection:
             return await connection.fetch_one(query=query)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> None:
+    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
         async with self.connection() as connection:
             return await connection.execute(query=query, values=values)
 
@@ -162,8 +162,8 @@ async def fetch_all(self, query: ClauseElement) -> typing.Any:
     async def fetch_one(self, query: ClauseElement) -> typing.Any:
         return await self._connection.fetch_one(query=query)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> None:
-        await self._connection.execute(query, values)
+    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+        return await self._connection.execute(query, values)
 
     async def execute_many(self, query: ClauseElement, values: list) -> None:
         await self._connection.execute_many(query, values)

tests/test_databases.py

@@ -151,6 +151,25 @@ async def test_queries(database_url):
             assert iterate_results[2]["completed"] == True
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_execute_return_val(database_url):
+    """
+    Test using return value from `execute()` to get an inserted primary key.
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            query = notes.insert()
+            values = {"text": "example1", "completed": True}
+            pk = await database.execute(query, values)
+
+            assert isinstance(pk, int)
+            query = notes.select().where(notes.c.id == pk)
+            result = await database.fetch_one(query)
+            assert result["text"] == "example1"
+            assert result["completed"] == True
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_rollback_isolation(database_url):