MagicStack · srimurthy211 · Jan 8, 2021
diff --git a/_python/pgbench_python.py b/_python/pgbench_python.py
@@ -42,7 +42,7 @@ def _ctr(_):
 
 def psycopg_connect(args):
     conn = psycopg2.connect(user=args.pguser, host=args.pghost,
-                            port=args.pgport)
+                            port=args.pgport, database=args.pgdb)
     return conn
 
 
@@ -67,7 +67,7 @@ def psycopg_copy(conn, query, args):
 
 def pypostgresql_connect(args):
     conn = postgresql.open(user=args.pguser, host=args.pghost,
-                           port=args.pgport)
+                           port=args.pgport, database=args.pgdb)
     return conn
 
 
@@ -78,7 +78,7 @@ def pypostgresql_execute(conn, query, args):
 
 async def aiopg_connect(args):
     conn = await aiopg.connect(user=args.pguser, host=args.pghost,
-                               port=args.pgport)
+                               port=args.pgport, database=args.pgdb)
     return conn
 
 
@@ -124,7 +124,7 @@ async def aiopg_tuples_executemany(conn, query, rows):
 
 async def asyncpg_connect(args):
     conn = await asyncpg.connect(user=args.pguser, host=args.pghost,
-                                 port=args.pgport)
+                                 port=args.pgport, database=args.pgdb)
     return conn
 
 
@@ -346,6 +346,8 @@ def die(msg):
     parser.add_argument(
         '-C', '--concurrency', type=int, default=10,
         help='number of concurrent connections')
+    parser.add_argument(
+        '--pgdb', '--database', type=str, help='database')
     parser.add_argument(
         '-D', '--duration', type=int, default=30,
         help='duration of test in seconds')

diff --git a/pgbench b/pgbench
@@ -365,7 +365,7 @@ BENCHMARKS = [
 ]
 
 
-def run_benchmarks(args, variations, pghost, pgport, pguser):
+def run_benchmarks(args, variations, pghost, pgport, pguser, pgdb):
     if args.benchmark:
         benchmarks_to_run = [re.compile(b) for b in args.benchmark]
     else:
@@ -394,6 +394,7 @@ def run_benchmarks(args, variations, pghost, pgport, pguser):
             'pghost': pghost,
             'pgport': pgport,
             'pguser': pguser,
+            'pgdb': pgdb,
             'output-format': 'json'
         }
 
@@ -477,6 +478,8 @@ def main():
     parser.add_argument(
         '--pguser', type=str, default='postgres',
         help='PostgreSQL server user')
+    parser.add_argument(
+        '--pgdb', type=str, help='database string')
     parser.add_argument(
         '--queryfiles', help='comma-separated list of files with queries')
     parser.add_argument(
@@ -533,13 +536,14 @@ def main():
         conn_spec = {
             'host': args.pghost,
             'port': args.pgport,
-            'user': args.pguser
+            'user': args.pguser,
+            'database': args.pgdb
         }
 
     try:
         benchmarks_data = run_benchmarks(
             args, variations, conn_spec['host'],
-            conn_spec['port'], conn_spec['user'])
+            conn_spec['port'], conn_spec['user'], conn_spec['database'])
     finally:
         if pg_cluster is not None:
             pg_cluster.stop()