Created
November 22, 2021 20:50
-
-
Save kmuthukk/2c4934c0f43a82297811e79f8d414e4a to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# Dependencies: | |
# On CentOS you can install psycopg2 thus: | |
# | |
# sudo yum install postgresql-libs | |
# sudo yum install python-psycopg2 | |
import psycopg2; | |
from multiprocessing.dummy import Pool as ThreadPool | |
num_tenants=100 | |
num_users=100 | |
num_threads=1 | |
connect_string="host=127.0.0.1 dbname=yugabyte user=yugabyte port=5433" | |
def create_table(): | |
conn = psycopg2.connect(connect_string) | |
conn.set_session(autocommit=True) | |
cur = conn.cursor() | |
cur.execute("""DROP TABLE IF EXISTS users"""); | |
cur.execute("""DROP TABLE IF EXISTS balance"""); | |
cur.execute("""CREATE TABLE IF NOT EXISTS users( | |
tenant_id text, | |
user_id text, | |
age int, | |
city text, | |
PRIMARY KEY(tenant_id, user_id)) | |
""") | |
cur.execute("""CREATE TABLE IF NOT EXISTS balance( | |
tenant_id text, | |
user_id text, | |
balance int, | |
PRIMARY KEY(tenant_id, user_id)) | |
""") | |
print("Created users & balance table") | |
print("====================") | |
def load_data_slave(thread_num): | |
thread_id = str(thread_num) | |
conn = psycopg2.connect(connect_string) | |
conn.set_session(autocommit=True) | |
cur = conn.cursor() | |
print("Thread-" + thread_id + ": Inserting %d rows..." % (num_users)) | |
num_inserts = 0 | |
try: | |
for idx in range(num_tenants): | |
for jdx in range(num_users): | |
cur.execute("""INSERT INTO users (tenant_id, user_id, age, city) VALUES (%s, %s, %s, %s)""", | |
("tenant-"+thread_id+"-"+str(idx), | |
"user--"+str(jdx), | |
20 + (jdx % 50), | |
"city--"+str(jdx % 1000))) | |
cur.execute("""INSERT INTO balance (tenant_id, user_id, balance) VALUES (%s, %s, %s)""", | |
("tenant-"+thread_id+"-"+str(idx), | |
"user--"+str(jdx), | |
20 + (jdx % 50))) | |
num_inserts += 1 | |
if (idx % 1000 == 0): | |
print("Thread-" + thread_id + ": Inserted %d rows" % (num_inserts)) | |
except Exception as e: | |
print("Unexpected exception: " + str(e)) | |
print("Thread-" + thread_id + ": Inserted %d rows" % (num_inserts)) | |
def load_data(): | |
pool = ThreadPool(num_threads) | |
results = pool.map(load_data_slave, range(num_threads)) | |
# Main | |
create_table() | |
load_data() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment