################################################################### # 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()