diff --git a/_python/pgbench_python.py b/_python/pgbench_python.py index 0616ce9..bb830a0 100755 --- 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 index d2d5d47..28991d2 100755 --- 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()