generate_tables.py

#!/usr/bin/python

import threading
import time
import os
import random

exitFlag = 0

class myThread (threading.Thread):
   def __init__(self, threadID):
      threading.Thread.__init__(self)
      self.threadID = threadID


   def run(self):
      commandList = []
      num_tables_per_tenant = random.randint(20,100)
      for i in range(0,num_tables_per_tenant):
         commandList.append("CREATE TABLE test_%s_%s (a int, b int);" % (self.threadID, i))

         # make the first table non-colocated with any table
         # rest all colocated with the first one
         if i == 0:
            commandList.append("SELECT create_distributed_table('test_%s_%s', 'a', colocate_with:='none');" % (self.threadID, i))
         else:
            commandList.append("SELECT create_distributed_table('test_%s_%s', 'a', colocate_with:='test_%s_0');" % (self.threadID, i, self.threadID))

      commandList.append("WITH shardids AS (SELECT shardid FROM pg_dist_shard ORDER BY shardid DESC LIMIT 1) SELECT citus_move_shard_placement(shardids.shardid, 'localhost', 9700, 'localhost', '9701', 'force_logical') FROM shardids;")
      
      for i in range(0,num_tables_per_tenant):
         commandList.append("SELECT count(*) FROM test_%s_%s;" % (self.threadID, i))

      run_command_on_psql(commandList)


def run_command_on_psql(cmdList):

   for cmd in cmdList:
      cmd_to_run = ("psql postgres -c \"%s\"" % cmd)
      print (cmd_to_run)
      ret = os.system(cmd_to_run)

      # detected failure in a command, exit
      if ret != 0:
         exit(-1)

# Create new threads
num_iterations = 1
for i in range(0, num_iterations):
   threads = []
   for n in range(0,4):
       thread_id = random.randint(0,10000000000)
       t = myThread(thread_id)
       t.start()
       threads.append(t)

   # Wait all threads to finish.
   for t in threads:
       t.join()


#thread2.start()