include psycopg3 benchmark (#9)

Author: farahats9

_python/pgbench_python.py

@@ -9,23 +9,23 @@
 
 import argparse
 import asyncio
-from concurrent import futures
 import csv
 import io
 import itertools
 import json
 import re
 import sys
 import time
-
-import numpy as np
-import uvloop
+from concurrent import futures
 
 import aiopg
 import asyncpg
+import numpy as np
 import postgresql
+import psycopg
 import psycopg2
 import psycopg2.extras
+import uvloop
 
 
 def _chunks(iterable, n):
@@ -41,18 +41,26 @@ def _ctr(_):
 
 
 def psycopg_connect(args):
-    conn = psycopg2.connect(user=args.pguser, host=args.pghost,
+    conn = psycopg.connect(user=args.pguser, host=args.pghost,
                             port=args.pgport)
     return conn
 
+def psycopg2_connect(args):
+    conn = psycopg2.connect(user=args.pguser, host=args.pghost,
+                            port=args.pgport)
+    return conn
 
-def psycopg_execute(conn, query, args):
+def psycopg2_execute(conn, query, args):
     cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
     cur.execute(query, args)
     return len(cur.fetchall())
 
+def psycopg_execute(conn, query, args):
+    cur = conn.cursor(row_factory=psycopg.rows.dict_row)
+    cur.execute(query, args)
+    return len(cur.fetchall())
 
-def psycopg_copy(conn, query, args):
+def psycopg2_copy(conn, query, args):
     rows, copy = args[:2]
     f = io.StringIO()
     writer = csv.writer(f, delimiter='\t')
@@ -64,6 +72,29 @@ def psycopg_copy(conn, query, args):
     conn.commit()
     return cur.rowcount
 
+def psycopg_copy(conn, query, args):
+    rows, copy = args[:2]
+    f = io.StringIO()
+    writer = csv.writer(f, delimiter='\t')
+    for row in rows:
+        writer.writerow(row)
+    f.seek(0)
+    with conn.cursor() as cur:
+        with cur.copy(query) as copy:
+            copy.write(f.read())
+        conn.commit()
+        return cur.rowcount
+
+def psycopg2_executemany(conn, query, args):
+    cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+    cur.executemany(query, args)
+    return len(args)
+
+def psycopg_executemany(conn, query, args):
+    with conn.cursor() as cur:
+        cur.executemany(query, args)
+    return len(args)
+
 
 def pypostgresql_connect(args):
     conn = postgresql.open(user=args.pguser, host=args.pghost,
@@ -128,15 +159,33 @@ async def asyncpg_connect(args):
     return conn
 
 
+async def async_psycopg_connect(args):
+    conn = await psycopg.AsyncConnection.connect(
+        user=args.pguser, host=args.pghost, port=args.pgport)
+    return conn
+
+
 async def asyncpg_execute(conn, query, args):
     return len(await conn.fetch(query, *args))
 
 
+async def async_psycopg_execute(conn, query, args):
+    cur = conn.cursor(row_factory=psycopg.rows.dict_row)
+    await cur.execute(query, args)
+    return len(await cur.fetchall())
+
+
 async def asyncpg_executemany(conn, query, args):
     await conn.executemany(query, args)
     return len(args)
 
 
+async def async_psycopg_executemany(conn, query, args):
+    async with conn.cursor() as cur:
+        await cur.executemany(query, args)
+    return len(args)
+
+
 async def asyncpg_copy(conn, query, args):
     rows, copy = args[:2]
     result = await conn.copy_records_to_table(
@@ -145,6 +194,21 @@ async def asyncpg_copy(conn, query, args):
     return int(count)
 
 
+async def async_psycopg_copy(conn, query, args):
+    rows, copy = args[:2]
+    f = io.StringIO()
+    writer = csv.writer(f, delimiter='\t')
+    for row in rows:
+        writer.writerow(row)
+    f.seek(0)
+    
+    async with conn.cursor() as cur:
+        async with cur.copy(query) as copy:
+            await copy.write(f.read())
+        await conn.commit()
+        return cur.rowcount
+
+
 async def worker(executor, eargs, start, duration, timeout):
     queries = 0
     rows = 0
@@ -369,7 +433,15 @@ def die(msg):
         help='PostgreSQL server user')
     parser.add_argument(
         'driver', help='driver implementation to use',
-        choices=['aiopg', 'aiopg-tuples', 'asyncpg', 'psycopg', 'postgresql'])
+        choices=[
+            'aiopg',
+            'aiopg-tuples',
+            'asyncpg',
+            'psycopg2',
+            'psycopg3',
+            'psycopg3-async',
+            'postgresql'
+        ],
     parser.add_argument(
         'queryfile', help='file to read benchmark query information from')
 
@@ -425,11 +497,25 @@ def die(msg):
             asyncpg_connect, asyncpg_execute, asyncpg_copy, asyncpg_executemany
         is_async = True
         arg_format = 'native'
-    elif args.driver == 'psycopg':
-        connector, executor, copy_executor = \
-            psycopg_connect, psycopg_execute, psycopg_copy
+    elif args.driver == 'psycopg2':
+        connector, executor, copy_executor, batch_executor = (
+            psycopg2_connect, psycopg2_execute, 
+            psycopg2_copy, psycopg2_executemany,
+        )
+        is_async = False
+        arg_format = 'python'
+    elif args.driver == 'psycopg3':
+        connector, executor, copy_executor, batch_executor = \
+            psycopg_connect, psycopg_execute, psycopg_copy, psycopg_executemany
         is_async = False
         arg_format = 'python'
+    elif args.driver == 'psycopg3-async':
+        connector, executor, copy_executor, batch_executor = (
+            async_psycopg_connect, async_psycopg_execute, 
+            async_psycopg_copy, async_psycopg_executemany,
+        )
+        is_async = True
+        arg_format = 'python'
     elif args.driver == 'postgresql':
         connector, executor = pypostgresql_connect, pypostgresql_execute
         is_async = False

_python/requirements.txt

@@ -1,6 +1,7 @@
 aiopg==1.0.0
-asyncpg==0.21.0
+asyncpg==0.27.0
 py-postgresql==1.2.2
-psycopg2-binary==2.8.6
+psycopg2-binary==2.9.6
 numpy~=1.19.0
 uvloop~=0.14.0
+psycopg[binary]==3.1.9

queries/7-oneplusone.json

@@ -0,0 +1,3 @@
+{
+    "query":"SELECT 1+1"
+}

requirements.txt

@@ -1,3 +1,3 @@
-asyncpg~=0.21.0
+asyncpg~=0.27.0
 distro~=1.5.0
 numpy~=1.19.0