autogen.py 4.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168
# -*- coding: utf-8 -*-

import sys
from util.log import *
from util.cases import *
from util.sql import *
import threading
import random
import string
import time


#
# Auto Gen class
#
class AutoGen:
    def __init__(self):
        self.ts = 1600000000000
        self.batch_size = 100
        seed = time.clock_gettime(time.CLOCK_REALTIME)
        random.seed(seed)

    # set start ts
    def set_start_ts(self, ts):
        self.ts = ts

    # set batch size
    def set_batch_size(self, batch_size):
        self.batch_size = batch_size

    #  _columns_sql
    def gen_columns_sql(self, pre, cnt, binary_len, nchar_len):
        types = [ 
            'timestamp',
            'tinyint',
            'smallint',
            'tinyint unsigned',
            'smallint unsigned',
            'int',
            'bigint',
            'int unsigned',
            'bigint unsigned',
            'float',
            'double',
            'bool',
            f'varchar({binary_len})',
            f'nchar({nchar_len})'
        ]

        sqls = ""
        metas = []
        for i in range(cnt):
            colname = f"{pre}{i}"
            if i < len(types):
               sel = i
            else:
               sel = random.randint(0, len(types)-1)
            coltype = types[sel]
            sql = f"{colname} {coltype}"
            if sqls != "":
                sqls += ","
            sqls += sql
            metas.append(sel)
        
        return metas, sqls;    

    # gen tags data
    def gen_data(self, i, marr):
        datas = ""   
        for c in marr:
            data = ""
            if c == 0 : # timestamp
                data = "%d" % (self.ts + i)
            elif c <= 4 : # small
                data = "%d"%(i%128)
            elif c <= 8 : # int
                data = f"{i}"
            elif c <= 10 : # float
                data = "%f"%(i+i/1000)
            elif c <= 11 : # bool
                data = "%d"%(i%2)
            elif c == 12 : # binary
                data = '"' + self.random_string(self.bin_len) + '"'
            elif c == 13 : # binary
                data = '"' + self.random_string(self.nch_len) + '"'

            if datas != "":
                datas += ","
            datas += data
        
        return datas

    # generate specail wide random string
    def random_string(self, count):
        letters = string.ascii_letters
        return ''.join(random.choice(letters) for i in range(count))

    # create db
    def create_db(self, dbname):
        self.dbname  = dbname
        tdSql.execute(f'create database {dbname}')
        tdSql.execute(f'use {dbname}')
        
    # create table or stable
    def create_stable(self, stbname, tag_cnt, column_cnt, binary_len, nchar_len):
        self.bin_len = binary_len
        self.nch_len = nchar_len
        self.stbname = stbname
        self.mtags, tags = self.gen_columns_sql("t", tag_cnt, binary_len, nchar_len)
        self.mcols, cols = self.gen_columns_sql("c", column_cnt - 1, binary_len, nchar_len)

        sql = f"create table {stbname} (ts timestamp, {cols}) tags({tags})"
        tdSql.execute(sql)

    # create child table 
    def create_child(self, stbname, prename, cnt):
        self.child_cnt = cnt
        self.child_name = prename
        for i in range(cnt):
            tags_data = self.gen_data(i, self.mtags)
            sql = f"create table {prename}{i} using {stbname} tags({tags_data})"
            tdSql.execute(sql)

        tdLog.info(f"create child tables {cnt} ok")

    def insert_data_child(self, child_name, cnt, batch_size, step):        
        values = ""
        print("insert child data")
        ts = self.ts

        # loop do
        for i in range(cnt):
            value = self.gen_data(i, self.mcols)
            ts += step
            values += f"({ts},{value}) "
            if batch_size == 1 or (i > 0 and i % batch_size == 0) :
                sql = f"insert into {child_name} values {values}"
                tdSql.execute(sql)
                if batch_size > 40:
                   tdLog.info(f" insert data i={i}")
                values = ""

        # end batch
        if values != "":
            sql = f"insert into {child_name} values {values}"
            tdSql.execute(sql)
            tdLog.info(f" insert data i={i}")
            values = ""

        tdLog.info(f" insert child data {child_name} finished, insert rows={cnt}")    

    # insert data 
    def insert_data(self, cnt):
        for i in range(self.child_cnt):
            name = f"{self.child_name}{i}"
            self.insert_data_child(name, cnt, self.batch_size, 1)

        tdLog.info(f" insert data ok, child table={self.child_cnt} insert rows={cnt}")

    # insert same timestamp to all childs
    def insert_samets(self, cnt):
        for i in range(self.child_cnt):
            name = f"{self.child_name}{i}"
            self.insert_data_child(name, cnt, self.batch_size, 0)

        tdLog.info(f" insert same timestamp ok, child table={self.child_cnt} insert rows={cnt}")