Script to creates lots of tables in a colocated database

ysql_many_colocated_tables.py

# 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
import time

host = "change-to-ip"
password = "change-to-password"

def drop_and_create_db(db_name):
  connect_string = "host={} dbname=yugabyte user=yugabyte password={} port=5433".format(host, password);
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor()
  print("dropping database if exists {}".format(db_name))
  cur.execute("DROP DATABASE IF EXISTS {}".format(db_name))
  print("dropped database if exists {}".format(db_name))
  cur.execute("CREATE DATABASE {} COLOCATED=true".format(db_name))
  print("created database {} COLOCATED=true".format(db_name))


def create_tables(dbname, num_tables):
  connect_string = "host={} dbname=mydb user=yugabyte password={} port=5433".format(host, password);
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor()
  for idx in range(num_tables):
    t1 = time.time();
    cur.execute(("CREATE TABLE IF NOT EXISTS t_{}"
                 + "(k int primary key, v1 text, v2 text)"
                ).format(idx))
    t2 = time.time();
    delta_ms = (t2 - t1) * 1000.0
    print("Created table t_{} in {} ms".format(idx, delta_ms))

# Main
num_tables = 60000
db_name = "mydb"
drop_and_create_db(db_name)
create_tables(db_name, num_tables)