tmqCommon.py 24.2 KB
Newer Older
P
plum-lihui 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
###################################################################
#           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 -*-

P
plum-lihui 已提交
14
from asyncore import loop
P
plum-lihui 已提交
15
from collections import defaultdict
P
plum-lihui 已提交
16
import subprocess
P
plum-lihui 已提交
17 18 19 20 21 22
import random
import string
import threading
import requests
import time
# import socketfrom
P
plum-lihui 已提交
23 24
import json
import toml
P
plum-lihui 已提交
25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48

import taos
from util.log import *
from util.sql import *
from util.cases import *
from util.dnodes import *
from util.common import *

# class actionType(Enum):
#     CREATE_DATABASE = 0
#     CREATE_STABLE   = 1
#     CREATE_CTABLE   = 2
#     INSERT_DATA     = 3

class TMQCom:
    def init(self, conn, logSql):
        tdSql.init(conn.cursor())
        # tdSql.init(conn.cursor(), logSql)  # output sql.txt file

    def initConsumerTable(self,cdbName='cdb'):        
        tdLog.info("create consume database, and consume info table, and consume result table")
        tdSql.query("create database if not exists %s vgroups 1"%(cdbName))
        tdSql.query("drop table if exists %s.consumeinfo "%(cdbName))
        tdSql.query("drop table if exists %s.consumeresult "%(cdbName))
P
plum-lihui 已提交
49
        tdSql.query("drop table if exists %s.notifyinfo "%(cdbName))      
P
plum-lihui 已提交
50 51 52

        tdSql.query("create table %s.consumeinfo (ts timestamp, consumerid int, topiclist binary(1024), keylist binary(1024), expectmsgcnt bigint, ifcheckdata int, ifmanualcommit int)"%cdbName)
        tdSql.query("create table %s.consumeresult (ts timestamp, consumerid int, consummsgcnt bigint, consumrowcnt bigint, checkresult int)"%cdbName)
P
plum-lihui 已提交
53
        tdSql.query("create table %s.notifyinfo (ts timestamp, cmdid int, consumerid int)"%cdbName)
P
plum-lihui 已提交
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

    def initConsumerInfoTable(self,cdbName='cdb'):        
        tdLog.info("drop consumeinfo table")
        tdSql.query("drop table if exists %s.consumeinfo "%(cdbName))
        tdSql.query("create table %s.consumeinfo (ts timestamp, consumerid int, topiclist binary(1024), keylist binary(1024), expectmsgcnt bigint, ifcheckdata int, ifmanualcommit int)"%cdbName)

    def insertConsumerInfo(self,consumerId, expectrowcnt,topicList,keyList,ifcheckdata,ifmanualcommit,cdbName='cdb'):    
        sql = "insert into %s.consumeinfo values "%cdbName
        sql += "(now, %d, '%s', '%s', %d, %d, %d)"%(consumerId, topicList, keyList, expectrowcnt, ifcheckdata, ifmanualcommit)
        tdLog.info("consume info sql: %s"%sql)
        tdSql.query(sql)

    def selectConsumeResult(self,expectRows,cdbName='cdb'):
        resultList=[]
        while 1:
            tdSql.query("select * from %s.consumeresult"%cdbName)
            #tdLog.info("row: %d, %l64d, %l64d"%(tdSql.getData(0, 1),tdSql.getData(0, 2),tdSql.getData(0, 3))
            if tdSql.getRows() == expectRows:
                break
            else:
                time.sleep(5)
        
        for i in range(expectRows):
            tdLog.info ("consume id: %d, consume msgs: %d, consume rows: %d"%(tdSql.getData(i , 1), tdSql.getData(i , 2), tdSql.getData(i , 3)))
            resultList.append(tdSql.getData(i , 3))
        
        return resultList
P
plum-lihui 已提交
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
    
    def selectConsumeMsgResult(self,expectRows,cdbName='cdb'):
        resultList=[]
        while 1:
            tdSql.query("select * from %s.consumeresult"%cdbName)
            #tdLog.info("row: %d, %l64d, %l64d"%(tdSql.getData(0, 1),tdSql.getData(0, 2),tdSql.getData(0, 3))
            if tdSql.getRows() == expectRows:
                break
            else:
                time.sleep(5)
        
        for i in range(expectRows):
            tdLog.info ("consume id: %d, consume msgs: %d, consume rows: %d"%(tdSql.getData(i , 1), tdSql.getData(i , 2), tdSql.getData(i , 3)))
            resultList.append(tdSql.getData(i , 2))
        
        return resultList
P
plum-lihui 已提交
97

P
plum-lihui 已提交
98
    def startTmqSimProcess(self,pollDelay,dbName,showMsg=1,showRow=1,cdbName='cdb',valgrind=0,alias=0,snapshot=0):
P
plum-lihui 已提交
99 100 101 102 103 104 105 106
        buildPath = tdCom.getBuildPath()
        cfgPath = tdCom.getClientCfgPath()
        if valgrind == 1:
            logFile = cfgPath + '/../log/valgrind-tmq.log'
            shellCmd = 'nohup valgrind --log-file=' + logFile
            shellCmd += '--tool=memcheck --leak-check=full --show-reachable=no --track-origins=yes --show-leak-kinds=all --num-callers=20 -v --workaround-gcc296-bugs=yes '
        
        if (platform.system().lower() == 'windows'):
wafwerar's avatar
wafwerar 已提交
107 108 109 110 111 112 113
            processorName = buildPath + '\\build\\bin\\tmq_sim.exe'
            if alias != 0:
                processorNameNew = buildPath + '\\build\\bin\\tmq_sim_new.exe'
                shellCmd = 'cp %s %s'%(processorName, processorNameNew)
                os.system(shellCmd)
                processorName = processorNameNew
            shellCmd = 'mintty -h never ' + processorName + ' -c ' + cfgPath
P
plum-lihui 已提交
114
            shellCmd += " -y %d -d %s -g %d -r %d -w %s -e %d "%(pollDelay, dbName, showMsg, showRow, cdbName, snapshot) 
P
plum-lihui 已提交
115 116
            shellCmd += "> nul 2>&1 &"   
        else:
P
plum-lihui 已提交
117 118 119 120 121 122 123
            processorName = buildPath + '/build/bin/tmq_sim'
            if alias != 0:
                processorNameNew = buildPath + '/build/bin/tmq_sim_new'
                shellCmd = 'cp %s %s'%(processorName, processorNameNew)
                os.system(shellCmd)
                processorName = processorNameNew
            shellCmd = 'nohup ' + processorName + ' -c ' + cfgPath
P
plum-lihui 已提交
124
            shellCmd += " -y %d -d %s -g %d -r %d -w %s -e %d "%(pollDelay, dbName, showMsg, showRow, cdbName, snapshot) 
P
plum-lihui 已提交
125 126
            shellCmd += "> /dev/null 2>&1 &"
        tdLog.info(shellCmd)
P
plum-lihui 已提交
127
        os.system(shellCmd) 
P
plum-lihui 已提交
128

P
plum-lihui 已提交
129 130 131 132 133 134 135 136 137 138 139 140 141
    def stopTmqSimProcess(self, processorName):
        psCmd = "ps -ef|grep -w %s|grep -v grep | awk '{print $2}'"%(processorName)
        processID = subprocess.check_output(psCmd, shell=True).decode("utf-8")
        while(processID):
            killCmd = "kill -INT %s > /dev/null 2>&1" % processID
            os.system(killCmd)
            time.sleep(0.2)
            processID = subprocess.check_output(psCmd, shell=True).decode("utf-8")
        tdLog.debug("%s is stopped by kill -INT" % (processorName))

    def getStartConsumeNotifyFromTmqsim(self,cdbName='cdb',rows=1):
        loopFlag = 1
        while loopFlag:
P
plum-lihui 已提交
142 143
            tdSql.query("select * from %s.notifyinfo"%cdbName)
            #tdLog.info("row: %d, %l64d, %l64d"%(tdSql.getData(0, 1),tdSql.getData(0, 2),tdSql.getData(0, 3))
P
plum-lihui 已提交
144 145 146 147 148 149 150
            actRows = tdSql.getRows()
            if (actRows >= rows):
                for i in range(actRows):
                    if tdSql.getData(i, 1) == 0:
                        loopFlag = 0
                        break            
            time.sleep(0.1)
P
plum-lihui 已提交
151 152
        return

P
plum-lihui 已提交
153 154 155
    def getStartCommitNotifyFromTmqsim(self,cdbName='cdb',rows=2):
        loopFlag = 1
        while loopFlag:
P
plum-lihui 已提交
156 157
            tdSql.query("select * from %s.notifyinfo"%cdbName)
            #tdLog.info("row: %d, %l64d, %l64d"%(tdSql.getData(0, 1),tdSql.getData(0, 2),tdSql.getData(0, 3))
P
plum-lihui 已提交
158 159 160 161 162 163
            actRows = tdSql.getRows()
            if (actRows >= rows):
                for i in range(actRows):
                    if tdSql.getData(i, 1) == 1:
                        loopFlag = 0
                        break            
P
plum-lihui 已提交
164 165 166
            time.sleep(0.1)
        return

P
plum-lihui 已提交
167 168 169 170 171 172 173 174
    def create_database(self,tsql, dbName,dropFlag=1,vgroups=4,replica=1):
        if dropFlag == 1:
            tsql.execute("drop database if exists %s"%(dbName))

        tsql.execute("create database if not exists %s vgroups %d replica %d"%(dbName, vgroups, replica))
        tdLog.debug("complete to create database %s"%(dbName))
        return

P
plum-lihui 已提交
175 176
    # self.create_stable() and self.create_ctable() and self.insert_data_interlaceByMultiTbl() : The three functions are matched
    # schema: (ts timestamp, c1 int, c2 bigint, c3 double, c4 binary(32), c5 nchar(32), c6 timestamp) tags (t1 int, t2 bigint, t3 double, t4 binary(32), t5 nchar(32))
P
plum-lihui 已提交
177
    def create_stable(self,tsql, dbName,stbName):
P
plum-lihui 已提交
178 179
        schemaString = "(ts timestamp, c1 int, c2 bigint, c3 double, c4 binary(32), c5 nchar(32), c6 timestamp) tags (t1 int, t2 bigint, t3 double, t4 binary(32), t5 nchar(32))"
        tsql.execute("create table if not exists %s.%s %s"%(dbName, stbName, schemaString))
P
plum-lihui 已提交
180 181 182
        tdLog.debug("complete to create %s.%s" %(dbName, stbName))
        return

P
plum-lihui 已提交
183
    def create_ctable(self,tsql=None, dbName='dbx',stbName='stb',ctbPrefix='ctb',ctbNum=1,ctbStartIdx=0):
P
plum-lihui 已提交
184
        # tsql.execute("use %s" %dbName)
P
plum-lihui 已提交
185 186 187
        pre_create = "create table"
        sql = pre_create
        #tdLog.debug("doing create one  stable %s and %d  child table in %s  ..." %(stbname, count ,dbname))
P
plum-lihui 已提交
188 189
        batchNum = 10
        tblBatched  = 0
P
plum-lihui 已提交
190
        for i in range(ctbNum):
P
plum-lihui 已提交
191
            tagBinaryValue = 'beijing'
P
plum-lihui 已提交
192
            if (i % 2 == 0):
P
plum-lihui 已提交
193
                tagBinaryValue = 'shanghai'
P
plum-lihui 已提交
194
            elif (i % 3 == 0):
P
plum-lihui 已提交
195
                tagBinaryValue = 'changsha'
P
plum-lihui 已提交
196
            
P
plum-lihui 已提交
197 198 199
            sql += " %s.%s%d using %s.%s tags(%d, %d, %d, '%s', '%s')"%(dbName,ctbPrefix,i+ctbStartIdx,dbName,stbName,i+ctbStartIdx,i+ctbStartIdx,i+ctbStartIdx,tagBinaryValue,tagBinaryValue)
            tblBatched += 1
            if (i == ctbNum-1 ) or (tblBatched == batchNum):
P
plum-lihui 已提交
200
                tsql.execute(sql)
P
plum-lihui 已提交
201
                tblBatched = 0
P
plum-lihui 已提交
202
                sql = pre_create
P
plum-lihui 已提交
203

P
plum-lihui 已提交
204 205 206
        if sql != pre_create:
            tsql.execute(sql)
        
P
plum-lihui 已提交
207 208
        tdLog.debug("complete to create %d child tables by %s.%s" %(ctbNum, dbName, stbName))
        return    
P
plum-lihui 已提交
209

P
plum-lihui 已提交
210 211
    def drop_ctable(self, tsql, dbname=None, count=1, default_ctbname_prefix="ctb",ctbStartIdx=0):
        for _ in range(count):
P
plum-lihui 已提交
212
            create_ctable_sql = f'drop table if exists {dbname}.{default_ctbname_prefix}{ctbStartIdx};'
P
plum-lihui 已提交
213 214 215 216
            ctbStartIdx += 1
            tdLog.info("drop ctb sql: %s"%create_ctable_sql)
            tsql.execute(create_ctable_sql)

P
plum-lihui 已提交
217
    # schema: (ts timestamp, c1 int, c2 binary(16))
P
plum-lihui 已提交
218
    def insert_data(self,tsql,dbName,stbName,ctbNum,rowsPerTbl,batchNum,startTs=None):
P
plum-lihui 已提交
219 220 221 222 223
        tdLog.debug("start to insert data ............")
        tsql.execute("use %s" %dbName)
        pre_insert = "insert into "
        sql = pre_insert

P
plum-lihui 已提交
224 225 226
        if startTs is None:
            t = time.time()
            startTs = int(round(t * 1000))
P
plum-lihui 已提交
227 228
        #tdLog.debug("doing insert data into stable:%s rows:%d ..."%(stbName, allRows))
        for i in range(ctbNum):
P
plum-lihui 已提交
229
            rowsBatched = 0
P
plum-lihui 已提交
230 231 232
            sql += " %s%d values "%(stbName,i)
            for j in range(rowsPerTbl):
                sql += "(%d, %d, 'tmqrow_%d') "%(startTs + j, j, j)
P
plum-lihui 已提交
233 234
                rowsBatched += 1
                if ((rowsBatched == batchNum) or (j == rowsPerTbl - 1)):
P
plum-lihui 已提交
235
                    tsql.execute(sql)
P
plum-lihui 已提交
236
                    rowsBatched = 0
P
plum-lihui 已提交
237 238 239 240 241 242 243 244 245
                    if j < rowsPerTbl - 1:
                        sql = "insert into %s%d values " %(stbName,i)
                    else:
                        sql = "insert into "
        #end sql
        if sql != pre_insert:
            #print("insert sql:%s"%sql)
            tsql.execute(sql)
        tdLog.debug("insert data ............ [OK]")
P
plum-lihui 已提交
246 247
        return        

P
plum-lihui 已提交
248
    # schema: (ts timestamp, c1 int, c2 int, c3 binary(16))
P
plum-lihui 已提交
249 250 251 252 253 254 255 256 257 258
    def insert_data_1(self,tsql,dbName,ctbPrefix,ctbNum,rowsPerTbl,batchNum,startTs):
        tdLog.debug("start to insert data ............")
        tsql.execute("use %s" %dbName)
        pre_insert = "insert into "
        sql = pre_insert

        t = time.time()
        startTs = int(round(t * 1000))
        #tdLog.debug("doing insert data into stable:%s rows:%d ..."%(stbName, allRows))
        for i in range(ctbNum):
P
plum-lihui 已提交
259
            rowsBatched = 0
P
plum-lihui 已提交
260 261 262 263 264 265
            sql += " %s%d values "%(ctbPrefix,i)
            for j in range(rowsPerTbl):
                if (j % 2 == 0):
                    sql += "(%d, %d, %d, 'tmqrow_%d') "%(startTs + j, j, j, j)
                else:
                    sql += "(%d, %d, %d, 'tmqrow_%d') "%(startTs + j, j, -j, j)
P
plum-lihui 已提交
266 267
                rowsBatched += 1
                if ((rowsBatched == batchNum) or (j == rowsPerTbl - 1)):
P
plum-lihui 已提交
268
                    tsql.execute(sql)
P
plum-lihui 已提交
269
                    rowsBatched = 0
P
plum-lihui 已提交
270 271 272 273 274 275 276 277 278 279 280
                    if j < rowsPerTbl - 1:
                        sql = "insert into %s%d values " %(ctbPrefix,i)
                    else:
                        sql = "insert into "
        #end sql
        if sql != pre_insert:
            #print("insert sql:%s"%sql)
            tsql.execute(sql)
        tdLog.debug("insert data ............ [OK]")
        return

P
plum-lihui 已提交
281
    # schema: (ts timestamp, c1 int, c2 int, c3 binary(16), c4 timestamp)
P
plum-lihui 已提交
282
    def insert_data_2(self,tsql,dbName,ctbPrefix,ctbNum,rowsPerTbl,batchNum,startTs,ctbStartIdx=0):
P
plum-lihui 已提交
283 284 285 286 287 288 289 290 291
        tdLog.debug("start to insert data ............")
        tsql.execute("use %s" %dbName)
        pre_insert = "insert into "
        sql = pre_insert

        t = time.time()
        startTs = int(round(t * 1000))
        #tdLog.debug("doing insert data into stable:%s rows:%d ..."%(stbName, allRows))
        for i in range(ctbNum):
P
plum-lihui 已提交
292
            rowsBatched = 0
P
plum-lihui 已提交
293
            sql += " %s%d values "%(ctbPrefix,i+ctbStartIdx)
P
plum-lihui 已提交
294 295 296 297 298
            for j in range(rowsPerTbl):
                if (j % 2 == 0):
                    sql += "(%d, %d, %d, 'tmqrow_%d', now) "%(startTs + j, j, j, j)
                else:
                    sql += "(%d, %d, %d, 'tmqrow_%d', now) "%(startTs + j, j, -j, j)
P
plum-lihui 已提交
299 300
                rowsBatched += 1
                if (rowsBatched == batchNum) or (j == rowsPerTbl - 1):
P
plum-lihui 已提交
301
                    tsql.execute(sql)
P
plum-lihui 已提交
302
                    rowsBatched = 0
P
plum-lihui 已提交
303
                    if j < rowsPerTbl - 1:
P
plum-lihui 已提交
304
                        sql = "insert into %s%d values " %(ctbPrefix,i+ctbStartIdx)
P
plum-lihui 已提交
305 306 307 308 309 310 311 312 313
                    else:
                        sql = "insert into "
        #end sql
        if sql != pre_insert:
            #print("insert sql:%s"%sql)
            tsql.execute(sql)
        tdLog.debug("insert data ............ [OK]")
        return

P
plum-lihui 已提交
314 315
    # schema: (ts timestamp, c1 int, c2 bigint, c3 double, c4 binary(32), c5 nchar(32), c6 timestamp) tags (t1 int, t2 bigint, t3 double, t4 binary(32), t5 nchar(32))
    def insert_data_interlaceByMultiTbl(self,tsql,dbName,ctbPrefix,ctbNum,rowsPerTbl,batchNum,startTs=0,ctbStartIdx=0):
P
plum-lihui 已提交
316 317 318 319 320 321 322 323 324 325 326 327 328 329
        tdLog.debug("start to insert data ............")
        tsql.execute("use %s" %dbName)
        pre_insert = "insert into "
        sql = pre_insert

        if startTs == 0:
            t = time.time()
            startTs = int(round(t * 1000))

        ctbDict = {}
        for i in range(ctbNum):
            ctbDict[i] = 0

        #tdLog.debug("doing insert data into stable:%s rows:%d ..."%(stbName, allRows))
P
plum-lihui 已提交
330
        rowsOfCtb = 0
P
plum-lihui 已提交
331 332
        while rowsOfCtb < rowsPerTbl:
            for i in range(ctbNum):
P
plum-lihui 已提交
333 334
                sql += " %s.%s%d values "%(dbName,ctbPrefix,i+ctbStartIdx)
                rowsBatched = 0
P
plum-lihui 已提交
335
                for k in range(batchNum):
P
plum-lihui 已提交
336
                    if (k % 2 == 0):
P
plum-lihui 已提交
337
                        sql += "(%d, %d, %d, %d, 'binary_%d', 'nchar_%d', now) "%(startTs+ctbDict[i], ctbDict[i],ctbDict[i], ctbDict[i],i+ctbStartIdx,k)
P
plum-lihui 已提交
338
                    else:
P
plum-lihui 已提交
339
                        sql += "(%d, %d, %d, %d, 'binary_%d', 'nchar_%d', now) "%(startTs+ctbDict[i],-ctbDict[i],ctbDict[i],-ctbDict[i],i+ctbStartIdx,k)
P
plum-lihui 已提交
340 341

                    rowsBatched += 1
P
plum-lihui 已提交
342
                    ctbDict[i] += 1
P
plum-lihui 已提交
343
                    if (rowsBatched == batchNum) or (ctbDict[i] == rowsPerTbl):
P
plum-lihui 已提交
344
                        tsql.execute(sql)
P
plum-lihui 已提交
345
                        rowsBatched = 0
P
plum-lihui 已提交
346 347 348 349 350 351 352
                        sql = "insert into "
                        break
            rowsOfCtb = ctbDict[0]

        tdLog.debug("insert data ............ [OK]")
        return

P
plum-lihui 已提交
353 354 355 356 357 358 359 360 361 362 363 364
    def threadFunctionForInsertByInterlace(self, **paraDict):
        # create new connector for new tdSql instance in my thread
        newTdSql = tdCom.newTdSql()
        self.insert_data_interlaceByMultiTbl(newTdSql,paraDict["dbName"],paraDict["ctbPrefix"],paraDict["ctbNum"],paraDict["rowsPerTbl"],paraDict["batchNum"],paraDict["startTs"],paraDict["ctbStartIdx"])
        return

    def asyncInsertDataByInterlace(self, paraDict):
        pThread = threading.Thread(target=self.threadFunctionForInsertByInterlace, kwargs=paraDict)
        pThread.start()
        return pThread

    def insert_data_with_autoCreateTbl(self,tsql,dbName,stbName,ctbPrefix,ctbNum,rowsPerTbl,batchNum,startTs=0,ctbStartIdx=0):
P
plum-lihui 已提交
365
        tdLog.debug("start to insert data with auto create child table ............")
P
plum-lihui 已提交
366 367 368 369 370 371 372
        tsql.execute("use %s" %dbName)
        pre_insert = "insert into "
        sql = pre_insert

        if startTs == 0:
            t = time.time()
            startTs = int(round(t * 1000))
P
plum-lihui 已提交
373
            
P
plum-lihui 已提交
374
        #tdLog.debug("doing insert data into stable:%s rows:%d ..."%(stbName, allRows))
P
plum-lihui 已提交
375
        rowsBatched = 0        
P
plum-lihui 已提交
376
        for i in range(ctbNum):
P
plum-lihui 已提交
377 378 379 380 381 382
            tagBinaryValue = 'beijing'
            if (i % 2 == 0):
                tagBinaryValue = 'shanghai'
            elif (i % 3 == 0):
                tagBinaryValue = 'changsha'
                
P
plum-lihui 已提交
383
            sql += " %s.%s%d using %s.%s tags (%d, %d, %d, '%s', '%s') values "%(dbName,ctbPrefix,i+ctbStartIdx,dbName,stbName,i+ctbStartIdx,i+ctbStartIdx,i+ctbStartIdx,tagBinaryValue,tagBinaryValue)
P
plum-lihui 已提交
384
            for j in range(rowsPerTbl):
P
plum-lihui 已提交
385
                sql += "(%d, %d, %d, %d, 'binary_%d', 'nchar_%d', now) "%(startTs+j, j,j, j,i+ctbStartIdx,rowsBatched)
P
plum-lihui 已提交
386 387
                rowsBatched += 1
                if ((rowsBatched == batchNum) or (j == rowsPerTbl - 1)):
P
plum-lihui 已提交
388
                    tsql.execute(sql)
P
plum-lihui 已提交
389
                    rowsBatched = 0
P
plum-lihui 已提交
390
                    if j < rowsPerTbl - 1:
P
plum-lihui 已提交
391
                        sql = "insert into %s.%s%d using %s.%s tags (%d, %d, %d, '%s', '%s') values " %(dbName,ctbPrefix,i+ctbStartIdx,dbName,stbName,i+ctbStartIdx,i+ctbStartIdx,i+ctbStartIdx,tagBinaryValue,tagBinaryValue)
P
plum-lihui 已提交
392 393 394 395 396 397 398
                    else:
                        sql = "insert into "
        #end sql
        if sql != pre_insert:
            #print("insert sql:%s"%sql)
            tsql.execute(sql)
        tdLog.debug("insert data ............ [OK]")
P
plum-lihui 已提交
399 400
        return

P
plum-lihui 已提交
401
    def syncCreateDbStbCtbInsertData(self, tsql, paraDict):
P
plum-lihui 已提交
402 403 404
        tdCom.create_database(tsql, paraDict["dbName"],paraDict["dropFlag"])
        tdCom.create_stable(tsql, dbname=paraDict["dbName"],stbname=paraDict["stbName"], column_elm_list=paraDict['colSchema'], tag_elm_list=paraDict['tagSchema'])
        tdCom.create_ctable(tsql, dbname=paraDict["dbName"],stbname=paraDict["stbName"],tag_elm_list=paraDict['tagSchema'],count=paraDict["ctbNum"], default_ctbname_prefix=paraDict['ctbPrefix'])
P
plum-lihui 已提交
405 406
        if "event" in paraDict and type(paraDict['event']) == type(threading.Event()):
            paraDict["event"].set()
P
plum-lihui 已提交
407 408 409 410 411 412

        ctbPrefix = paraDict['ctbPrefix']
        ctbNum = paraDict["ctbNum"]
        for i in range(ctbNum):
            tbName = '%s%s'%(ctbPrefix,i)
            tdCom.insert_rows(tsql,dbname=paraDict["dbName"],tbname=tbName,start_ts_value=paraDict['startTs'],count=paraDict['rowsPerTbl'])
P
plum-lihui 已提交
413 414 415 416 417 418 419 420 421 422 423 424 425
        return 

    def threadFunction(self, **paraDict):
        # create new connector for new tdSql instance in my thread
        newTdSql = tdCom.newTdSql()
        self.syncCreateDbStbCtbInsertData(self, newTdSql, paraDict)
        return

    def asyncCreateDbStbCtbInsertData(self, paraDict):
        pThread = threading.Thread(target=self.threadFunction, kwargs=paraDict)
        pThread.start()
        return pThread

P
plum-lihui 已提交
426 427 428
    def threadFunctionForInsert(self, **paraDict):
        # create new connector for new tdSql instance in my thread
        newTdSql = tdCom.newTdSql()
P
plum-lihui 已提交
429 430 431 432
        if 'ctbStartIdx' in paraDict.keys():
            self.insert_data_2(newTdSql,paraDict["dbName"],paraDict["ctbPrefix"],paraDict["ctbNum"],paraDict["rowsPerTbl"],paraDict["batchNum"],paraDict["startTs"],paraDict["ctbStartIdx"])
        else:
            self.insert_data_2(newTdSql,paraDict["dbName"],paraDict["ctbPrefix"],paraDict["ctbNum"],paraDict["rowsPerTbl"],paraDict["batchNum"],paraDict["startTs"])
P
plum-lihui 已提交
433 434 435 436 437 438 439
        return

    def asyncInsertData(self, paraDict):
        pThread = threading.Thread(target=self.threadFunctionForInsert, kwargs=paraDict)
        pThread.start()
        return pThread

P
plum-lihui 已提交
440
    def checkFileContent(self, consumerId, queryString, skipRowsOfCons=0):
P
plum-lihui 已提交
441 442 443 444 445 446 447 448 449 450 451 452 453 454 455
        buildPath = tdCom.getBuildPath()
        cfgPath = tdCom.getClientCfgPath()
        dstFile = '%s/../log/dstrows_%d.txt'%(cfgPath, consumerId)
        cmdStr = '%s/build/bin/taos -c %s -s "%s >> %s"'%(buildPath, cfgPath, queryString, dstFile)
        tdLog.info(cmdStr)
        os.system(cmdStr)
        
        consumeRowsFile = '%s/../log/consumerid_%d.txt'%(cfgPath, consumerId)
        tdLog.info("rows file: %s, %s"%(consumeRowsFile, dstFile))

        consumeFile = open(consumeRowsFile, mode='r')
        queryFile = open(dstFile, mode='r')
        
        # skip first line for it is schema
        queryFile.readline()
P
plum-lihui 已提交
456
        
P
plum-lihui 已提交
457 458 459 460
        # skip offset for consumer
        for i in range(0,skipRowsOfCons):
            consumeFile.readline()            
        
P
plum-lihui 已提交
461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507
        lines = 0
        while True:
            dst = queryFile.readline()
            src = consumeFile.readline()
            lines += 1
            if dst:
                if dst != src:
                    tdLog.info("src row: %s"%src)
                    tdLog.info("dst row: %s"%dst)
                    tdLog.exit("consumerId %d consume rows[%d] is not match the rows by direct query"%(consumerId, lines))
            else:
                break
        return 

    def getResultFileByTaosShell(self, consumerId, queryString):
        buildPath = tdCom.getBuildPath()
        cfgPath = tdCom.getClientCfgPath()
        dstFile = '%s/../log/dstrows_%d.txt'%(cfgPath, consumerId)
        cmdStr = '%s/build/bin/taos -c %s -s "%s >> %s"'%(buildPath, cfgPath, queryString, dstFile)
        tdLog.info(cmdStr)
        os.system(cmdStr)
        return dstFile
    
    def checkTmqConsumeFileContent(self, consumerId, dstFile):   
        cfgPath = tdCom.getClientCfgPath()     
        consumeRowsFile = '%s/../log/consumerid_%d.txt'%(cfgPath, consumerId)
        tdLog.info("rows file: %s, %s"%(consumeRowsFile, dstFile))

        consumeFile = open(consumeRowsFile, mode='r')
        queryFile = open(dstFile, mode='r')
        
        # skip first line for it is schema
        queryFile.readline()
        lines = 0
        while True:
            dst = queryFile.readline()
            src = consumeFile.readline()
            lines += 1
            if dst:
                if dst != src:
                    tdLog.info("src row: %s"%src)
                    tdLog.info("dst row: %s"%dst)
                    tdLog.exit("consumerId %d consume rows[%d] is not match the rows by direct query"%(consumerId, lines))
            else:
                break
        return 

P
plum-lihui 已提交
508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538
    def create_ntable(self, tsql, dbname=None, tbname_prefix="ntb", tbname_index_start_num = 1, column_elm_list=None, colPrefix='c', tblNum=1, **kwargs):
        tb_params = ""
        if len(kwargs) > 0:
            for param, value in kwargs.items():
                tb_params += f'{param} "{value}" '
        column_type_str = tdCom.gen_column_type_str(colPrefix, column_elm_list)

        for _ in range(tblNum):
            create_table_sql = f'create table {dbname}.{tbname_prefix}{tbname_index_start_num} ({column_type_str}) {tb_params};'
            tbname_index_start_num += 1
            tsql.execute(create_table_sql)

    def insert_rows_into_ntbl(self, tsql, dbname=None, tbname_prefix="ntb", tbname_index_start_num = 1, column_ele_list=None, startTs=None, tblNum=1, rows=1):
        if startTs is None:
            startTs = tdCom.genTs()[0]

        for tblIdx in range(tblNum):
            for rowIdx in range(rows):
                column_value_list = tdCom.gen_column_value_list(column_ele_list, f'{startTs}+{rowIdx}s')
                column_value_str = ''
                idx = 0
                for column_value in column_value_list:
                    if isinstance(column_value, str) and idx != 0:
                        column_value_str += f'"{column_value}", '
                    else:
                        column_value_str += f'{column_value}, '
                        idx += 1
                column_value_str = column_value_str.rstrip()[:-1]
                insert_sql = f'insert into {dbname}.{tbname_prefix}{tblIdx+tbname_index_start_num} values ({column_value_str});'
                tsql.execute(insert_sql)

P
plum-lihui 已提交
539 540 541 542
    def close(self):
        self.cursor.close()

tmqCom = TMQCom()