restfulInsert.py

###################################################################
#           Copyright (c) 2016 by TAOS Technologies, Inc.
#                     All rights reserved.
#
#  This file is proprietary and confidential to TAOS Technologies.
#  No part of this file may be reproduced, stored, transmitted,
#  disclosed or used in any form or by any means other than as
#  expressly provided by the written permission from Jianhui Tao
#
###################################################################

# -*- coding: utf-8 -*-

import requests
import threading
import random
import time
import argparse

class RestfulInsert:
    def __init__(self, host, startTimestamp, dbname, threads, tables, records, batchSize, tbNamePerfix, outOfOrder,tablePerbatch):
        self.header = {'Authorization': 'Basic cm9vdDp0YW9zZGF0YQ=='}
        self.url = "http://%s:6041/rest/sql" % host
        self.ts = startTimestamp
        self.dbname = dbname
        self.numOfThreads = threads
        self.numOfTables = tables
        self.recordsPerTable = records
        self.batchSize = batchSize
        self.tableNamePerfix = tbNamePerfix
        self.outOfOrder = outOfOrder
        self.tablePerbatch = tablePerbatch
    
    def createTable(self, threadID):
        tablesPerThread = int (self.numOfTables / self.numOfThreads)
        loop = tablesPerThread if threadID != self.numOfThreads - 1 else self.numOfTables - tablesPerThread * threadID        
        print("create table %d to %d" % (tablesPerThread * threadID, tablesPerThread * threadID + loop - 1))
        for i in range(loop):
            tableID = threadID * tablesPerThread
            if tableID + i >= self.numOfTables : break
            name = 'beijing' if (tableID + i) % 2 == 0 else 'shanghai'
            data = "create table if not exists %s.%s%d using %s.meters tags(%d, '%s')" % (self.dbname, self.tableNamePerfix, tableID + i, self.dbname, tableID + i, name)
            try:
                response = requests.post(self.url, data, headers = self.header)
                if response.status_code != 200:
                    print(response.content)                
            except Exception as e:
                print(e)

    def insertData(self, threadID):        
        print("thread %d started" % threadID)
        tablesPerThread = int (self.numOfTables / self.numOfThreads)   
        loop = int(self.recordsPerTable / self.batchSize)   
        if self.tablePerbatch == 1 : 
            for i in range(tablesPerThread+1):            
                tableID = i + threadID * tablesPerThread
                if tableID >= self.numOfTables: return
                start = self.ts
                start1=time.time()
                for k in range(loop):
                    data = "insert into %s.%s%d values" % (self.dbname, self.tableNamePerfix, tableID)
                    values = []
                    bloop = self.batchSize if k != loop - 1 else self.recordsPerTable - self.batchSize * k
                    for l in range(bloop):
                        values.append("(%d, %d, %d, %d)" %  (start + k * self.batchSize + l, random.randint(1, 100), random.randint(1, 100), random.randint(1, 100)))                              
                    if len(data) > 1048576 : 
                        print ('batch size is larger than 1M')
                        exit(-1)
                    if self.outOfOrder :
                        random.shuffle(values)
                    data+=''.join(values)
                    response = requests.post(self.url, data, headers = self.header)
                    if response.status_code != 200:
                        print(response.content)
        else:
            for i in range(0,tablesPerThread+self.tablePerbatch,self.tablePerbatch): 
                for k in range(loop):
                    data = "insert into "
                    for j in range(self.tablePerbatch):
                        tableID = i + threadID * tablesPerThread+j
                        if tableID >= self.numOfTables: return
                        start = self.ts
                        data += "%s.%s%d values" % (self.dbname, self.tableNamePerfix, tableID)
                        values = []
                        bloop = self.batchSize if k != loop - 1 else self.recordsPerTable - self.batchSize * k
                        for l in range(bloop):
                            values.append("(%d, %d, %d, %d)" %  (start + k * self.batchSize + l, random.randint(1, 100), random.randint(1, 100), random.randint(1, 100)))    
                        if self.outOfOrder :
                            random.shuffle(values)
                        data+=''.join(values)                          
                    if len(data) > 1024*1024 : 
                        print ('batch size is larger than 1M')
                        exit(-1)
                    try:
                        startTime = time.time()
                        response = requests.post(self.url, data, headers = self.header)
                        endTime = time.time()
                        if response.status_code != 200:
                            print(response.content)
                        else:
                            print("inserted %d records, %d seconds" % (bloop, endTime - startTime))
                    except Exception as e:
                        print(e)

    def insertUnlimitedData(self, threadID):        
        print("thread %d started" % threadID)
        tablesPerThread = int (self.numOfTables / self.numOfThreads)
        
        count = 0
        while True:
            i = 0
            start = self.ts  + count * self.batchSize          
            count = count + 1            
            
            for i in range(tablesPerThread):
                tableID = i + threadID * tablesPerThread
                                
                data = "insert into %s.%s%d values" % (self.dbname, self.tableNamePerfix, tableID)
                values = []
                for k in range(self.batchSize):
                    values.append("(%d, %d, %d, %d)" %  (start + k, random.randint(1, 100), random.randint(1, 100), random.randint(1, 100)))
                
                if(self.outOfOrder == False):
                    for k in range(len(values)):            
                        data += values[k]
                else:
                    random.shuffle(values)
                    for k in range(len(values)):            
                        data += values[k]    
                try:     
                    startTime = time.time()
                    response = requests.post(self.url, data, headers = self.header)
                    endTime = time.time()
                    if response.status_code != 200:
                        print(response.content)
                    else:
                        print("inserted %d records, %d seconds" % (self.batchSize, endTime - startTime))
                except Exception as e:
                    print(e)                    

    def run(self):            
        data = "create database if not exists %s" % self.dbname
        requests.post(self.url, data, headers = self.header)
        data = "create table if not exists %s.meters(ts timestamp, f1 int, f2 int, f3 int) tags(id int, loc nchar(20))" % self.dbname
        requests.post(self.url, data, headers = self.header)

        threads = []
        startTime = time.time()    
        for i in range(self.numOfThreads):
            thread = threading.Thread(target=self.createTable, args=(i,))
            thread.start()
            threads.append(thread)
        for i in range(self.numOfThreads):
            threads[i].join()
        print("createing %d tables takes %d seconds" % (self.numOfTables, (time.time() - startTime)))

        print("inserting data =======")
        threads = []
        startTime = time.time()
        for i in range(self.numOfThreads):
            if(self.recordsPerTable != -1):          
                thread = threading.Thread(target=self.insertData, args=(i,))
            else:
                thread = threading.Thread(target=self.insertUnlimitedData, args=(i,))
            thread.start()
            threads.append(thread)
        
        for i in range(self.numOfThreads):
            threads[i].join()
        print("inserting %s records takes %d seconds" % (self.numOfTables * self.recordsPerTable, (time.time() - startTime)))

parser = argparse.ArgumentParser()
parser.add_argument(
    '-H',
    '--host-name',
    action='store',
    default='127.0.0.1',
    type=str,
    help='host name to be connected (default: 127.0.0.1)')
parser.add_argument(
    '-S',
    '--start-timestamp',
    action='store',
    default=1500000000000,
    type=int,
    help='insert data from timestamp (default: 1500000000000)')
parser.add_argument(
    '-d',
    '--db-name',
    action='store',
    default='test',
    type=str,
    help='Database name to be created (default: test)')
parser.add_argument(
    '-t',
    '--number-of-threads',
    action='store',
    default=10,
    type=int,
    help='Number of threads to create tables and insert datas (default: 10)')
parser.add_argument(
    '-T',
    '--number-of-tables',
    action='store',
    default=10000,
    type=int,
    help='Number of tables to be created (default: 1000)')
parser.add_argument(
    '-r',
    '--number-of-records',
    action='store',
    default=10000,
    type=int,
    help='Number of record to be created for each table  (default: 1000, -1 for unlimited records)')
parser.add_argument(
    '-s',
    '--batch-size',
    action='store',
    default='1000',
    type=int,
    help='Number of tables to be created (default: 1000)')
parser.add_argument(
    '-p',
    '--table-name-prefix',
    action='store',
    default='t',
    type=str,
    help='Number of tables to be created (default: 1000)')
parser.add_argument(
    '-o',
    '--out-of-order',
    action='store_true', 
    help='The order of test data (default: False)')
parser.add_argument(
    '-b',
    '--table-per-batch',
    action='store', 
    default=1,
    type=int,
    help='the table per batch (default: 1)')



args = parser.parse_args()
ri = RestfulInsert(
        args.host_name, args.start_timestamp, args.db_name, args.number_of_threads, args.number_of_tables, 
        args.number_of_records, args.batch_size, args.table_name_prefix, args.out_of_order, args.table_per_batch)
ri.run()