Use psycopg2.extras.execute_batch instead of looping

elprans · elprans · commit 1b361556b422 · 2020-11-29T15:34:14.000-08:00
For fairness, psycopg2 has a faster "executemany" method:
`execute_batch`, and although it's not exposed via aiopg, it's easy to
reimplement.
diff --git a/.gitignore b/.gitignore
@@ -9,4 +9,5 @@ __pycache__/
 /*.egg-info
 /dist
 /.cache
+/.vscode
 docs/_build
diff --git a/_python/pgbench_python.py b/_python/pgbench_python.py
@@ -12,6 +12,7 @@
 from concurrent import futures
 import csv
 import io
+import itertools
 import json
 import re
 import sys
@@ -27,6 +28,18 @@
 import psycopg2.extras
 
 
+def _chunks(iterable, n):
+    i = 0
+
+    def _ctr(_):
+        nonlocal i
+        k = i // n
+        i += 1
+        return k
+    for _, g in itertools.groupby(iterable, _ctr):
+        yield g
+
+
 def psycopg_connect(args):
     conn = psycopg2.connect(user=args.pguser, host=args.pghost,
                             port=args.pgport)
@@ -77,12 +90,16 @@ async def aiopg_execute(conn, query, args):
     return rv
 
 
+async def _aiopg_executemany(cursor, query, rows):
+    for batch in _chunks(rows, n=100):
+        sqls = [cursor.mogrify(query, args) for args in batch]
+        await cursor.execute(b";".join(sqls))
+    return len(rows)
+
+
 async def aiopg_executemany(conn, query, rows):
     cur = await conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    rv = 0
-    for args in rows:
-        await cur.execute(query, args)
-        rv += cur.rowcount
+    rv = await _aiopg_executemany(cur, query, rows)
     cur.close()
     return rv
 
@@ -100,10 +117,7 @@ async def aiopg_tuples_execute(conn, query, args):
 
 async def aiopg_tuples_executemany(conn, query, rows):
     cur = await conn.cursor()
-    rv = 0
-    for args in rows:
-        await cur.execute(query, args)
-        rv += cur.rowcount
+    rv = await _aiopg_executemany(cur, query, rows)
     cur.close()
     return rv
 
