Update Connection.execute to accept query arguments

Author: 1st1

asyncpg/connection.py

@@ -112,12 +112,38 @@ def transaction(self, *, isolation='read_committed', readonly=False,
         """
         return transaction.Transaction(self, isolation, readonly, deferrable)
 
-    async def execute(self, script: str, *, timeout: float=None) -> str:
+    async def execute(self, query: str, *args, timeout: float=None) -> str:
         """Execute an SQL command (or commands).
 
+        This method can execute many SQL commands at once, when no arguments
+        are provided.
+
+        Example:
+
+        .. code-block:: pycon
+
+            >>> await con.execute('''
+            ...     CREATE TABLE mytab (a int);
+            ... ''')
+            CREATE TABLE
+
+            >>> await con.execute('''
+            ...     INSERT INTO mytab (a) VALUES ($1), ($2)
+            ... ''', 10, 20)
+            INSERT 0 2
+
+        :param args: Query arguments.
+        :param float timeout: Optional timeout value in seconds.
         :return str: Status of the last SQL command.
         """
-        return await self._protocol.query(script, timeout)
+        if not args:
+            return await self._protocol.query(query, timeout)
+
+        stmt = await self._get_statement(query, timeout)
+        protocol = self._protocol
+        _, status, _ = await protocol.bind_execute(stmt, args, '', 0,
+                                                   True, timeout)
+        return status.decode()
 
     async def _get_statement(self, query, timeout):
         cache = self._stmt_cache_max_size > 0

tests/test_execute.py

@@ -25,6 +25,29 @@ async def test_execute_script_1(self):
         self.assertEqual(self.con._protocol.queries_count, 1)
         self.assertEqual(status, 'SELECT 10')
 
+    async def test_execute_script_2(self):
+        status = await self.con.execute('''
+            CREATE TABLE mytab (a int);
+        ''')
+        self.assertEqual(status, 'CREATE TABLE')
+
+        try:
+            status = await self.con.execute('''
+                INSERT INTO mytab (a) VALUES ($1), ($2)
+            ''', 10, 20)
+            self.assertEqual(status, 'INSERT 0 2')
+        finally:
+            await self.con.execute('DROP TABLE mytab')
+
+    async def test_execute_script_3(self):
+        with self.assertRaisesRegex(asyncpg.PostgresSyntaxError,
+                                    'cannot insert multiple commands'):
+
+            await self.con.execute('''
+                CREATE TABLE mytab (a int);
+                INSERT INTO mytab (a) VALUES ($1), ($2);
+            ''', 10, 20)
+
     async def test_execute_script_check_transactionality(self):
         with self.assertRaises(asyncpg.PostgresError):
             await self.con.execute('''