【好文推荐openGauss psycopg2 批量操作性能测试

2023年 9月 12日 106.0k 0

# coding: utf-8

# Usage: python test.py

import time
import sys
import psycopg2 as pg
from psycopg2.extras import execute_batch, execute_values
from contextlib import contextmanager

if sys.argv[3] == "insert":
args = [[str(i), i] for i in range(int(sys.argv[2]))]
elif sys.argv[3] == "update":
args = [[i, str(i)] for i in range(int(sys.argv[2]))]
elif sys.argv[3] == "delete":
args = [[i] for i in range(int(sys.argv[2]))]
'''
- *dbname*: the database name
- *database*: the database name (only as keyword argument)
- *user*: user name used to authenticate
- *password*: password used to authenticate
- *host*: database host address (defaults to UNIX socket if not provided)
- *port*: connection port number (defaults to 5432 if not provided)
'''
conf = {
'dbname': "postgres",
'user': 'gaussdb',
'password': '',
'host': '',
'port': 26000,
'sslmode': 'disable'
}

@contextmanager
def calc_time(s):
start = time.time()
yield
end = time.time()
print(f"{s} of '{sys.argv[3]}' cost: ", end - start)

sql_map = {
"insert": {
1: "INSERT INTO t_psycopg2_benchmark VALUES (%s, %s)",
2: "INSERT INTO t_psycopg2_benchmark VALUES ($1, $2)",
3: "INSERT INTO t_psycopg2_benchmark VALUES %s",
},
"update": {
1: "UPDATE t_psycopg2_benchmark as t SET f_value = %s WHERE t.f_key = %s",
2: "UPDATE t_psycopg2_benchmark as t SET f_value = $1 WHERE t.f_key = $2",
3: "UPDATE t_psycopg2_benchmark as t SET f_value = data.v1 FROM (VALUES %s) AS data (id, v1) WHERE t.f_key = data.id",
},
"delete": {
1: "DELETE FROM t_psycopg2_benchmark as t WHERE t.f_key=%s",
2: "DELETE FROM t_psycopg2_benchmark as t WHERE t.f_key=$1",
3: "DELETE FROM t_psycopg2_benchmark as t WHERE t.f_key IN (%s)",
}
}

def insert_data(conn):
print("* preparing data ...")
args = [[str(i), i] for i in range(int(sys.argv[2]))]
cursor = conn.cursor()
sql = "insert into t_psycopg2_benchmark values %s"
execute_values(cursor, sql, args)
conn.commit()

def main():
try:
conn = pg.connect(**conf)
print("* connect success")
except Exception as e:
print(f"connect failed: {e}")
return

cursor = conn.cursor()

sql = "drop table if exists t_psycopg2_benchmark"
cursor.execute(sql)
sql = "create table t_psycopg2_benchmark (f_key text primary key, f_value numeric)"
cursor.execute(sql)

api = sys.argv[1]
if sys.argv[3] != "insert":
insert_data(conn)

print("* benchmarking ...")
if api == "executemany":
with calc_time("executemany"):
sql = sql_map[sys.argv[3]][1]
cursor.executemany(sql, args)
conn.commit()
elif api == "execute_batch":
with calc_time("execute_batch"):
sql = sql_map[sys.argv[3]][1]
execute_batch(cursor, sql, args)
conn.commit()
elif api == "prepare":
with calc_time("execute_values"):
cursor.execute(f"PREPARE test_stmt AS {sql_map[sys.argv[3]][2]}")
if sys.argv[3] == "delete":
execute_batch(cursor, "EXECUTE test_stmt (%s)", args)
else:
execute_batch(cursor, "EXECUTE test_stmt (%s, %s)", args)
cursor.execute("DEALLOCATE test_stmt")
conn.commit()
elif api == "execute_values":
with calc_time("execute_values"):
sql = sql_map[sys.argv[3]][3]
execute_values(cursor, sql, args)
conn.commit()
else:
print(f"unknow api: {api}")

if sys.argv[3] != "delete":
cursor.execute("delete from t_psycopg2_benchmark")
conn.commit()

if __name__ == "__main__":
main()

相关文章

Oracle如何使用授予和撤销权限的语法和示例
Awesome Project: 探索 MatrixOrigin 云原生分布式数据库
下载丨66页PDF,云和恩墨技术通讯(2024年7月刊)
社区版oceanbase安装
Oracle 导出CSV工具-sqluldr2
ETL数据集成丨快速将MySQL数据迁移至Doris数据库

发布评论