Elgato_dark/TDengine
1
0
Fork 0
mirror of https://github.com/taosdata/TDengine synced 2026-05-24 10:09:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4
TDengine/test/new_test_framework/utils/sql.py

3096 lines
113 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

###################################################################
# 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 sys
import os
import time
import datetime
import inspect
import traceback
import psutil
import shutil
import re
import pandas as pd
import csv
from .log import *
from .constant import *
import ctypes
import random
import datetime
import time
import taos
import platform
from tzlocal import get_localzone
from typing import Optional, Literal
def _parse_ns_timestamp(timestr):
dt_obj = datetime.datetime.strptime(
timestr[: len(timestr) - 3], "%Y-%m-%d %H:%M:%S.%f"
)
tz = (
int(
int(
(
dt_obj - datetime.datetime.fromtimestamp(0, dt_obj.tzinfo)
).total_seconds()
)
* 1e9
)
+ int(dt_obj.microsecond * 1000)
+ int(timestr[-3:])
)
return tz
def _parse_datetime(timestr):
"""
Parse a string to a datetime object. The string can be in one of the following formats:
The string can be in one of the following formats:
- '%Y-%m-%d %H:%M:%S.%f%z': Contains microseconds and timezone offset.
- '%Y-%m-%d %H:%M:%S%z': Contains no microseconds but contains timezone offset.
- '%Y-%m-%d %H:%M:%S.%f': Contains microseconds but no timezone offset.
- '%Y-%m-%d %H:%M:%S': Contains no microseconds and no timezone offset.
Args:
timestr (str): The string to be parsed.
Returns:
datetime.datetime: The datetime object parsed from the string.
"""
formats = [
"%Y-%m-%d %H:%M:%S.%f%z", # 包含微秒和时区偏移
"%Y-%m-%d %H:%M:%S%z", # 不包含微秒但包含时区偏移
"%Y-%m-%d %H:%M:%S.%f", # 包含微秒
"%Y-%m-%d %H:%M:%S", # 不包含微秒
]
for fmt in formats:
try:
# try to parse the string with the current format
dt = datetime.datetime.strptime(timestr, fmt)
# 如果字符串包含时区信息，则返回 aware 对象
# if sting contains timezone info, return aware object
if dt.tzinfo is not None:
return dt
else:
# if sting does not contain timezone info, assume it is in local timezone
# get local timezone
local_timezone = get_localzone()
# print("Timezone:", local_timezone)
return dt.replace(tzinfo=local_timezone)
except ValueError:
continue # if the current format does not match, try the next format
# 如果所有格式都不匹配，返回 None
# if none of the formats match, return
raise ValueError(
f"input format does not match. correct formats include: '{', '.join(formats)}'"
)
def _fast_caller(depth: int = 1):
try:
f = sys._getframe(depth)
return f.f_code.co_filename, f.f_lineno
except Exception:
fr = inspect.getframeinfo(inspect.stack()[depth][0])
return fr.filename, fr.lineno
class TDSql:
def __init__(self):
self.queryRows = 0
self.queryCols = 0
self.affectedRows = 0
self.csvLine = 0
self.replica = 1
def init(self, cursor, log=False):
"""
Initializes the TDSql instance with a database cursor and optionally enables logging.
Args:
cursor: The database cursor to be used for executing SQL queries.
log (bool, optional): If True, enables logging of SQL statements to a file. Defaults to False.
Returns:
None
Raises:
None
"""
self.cursor = cursor
self.sql = None
if log:
filename, _ = _fast_caller(1)
self.cursor.log(filename + ".sql")
def close(self):
"""
Closes the cursor.
Args:
None
Returns:
None
Raises:
None
"""
self.cursor.close()
def connect(self, username="root", passwd="taosdata", **kwargs):
"""
Reconnect
Args:
username (str):The username used to log in to the cluster.
passwd (str, optional): The password used to log in to the cluster.
Returns:
None
Raises:
None
"""
self.cursor.close()
if not kwargs:
kwargs = {
'user': username,
'password': passwd,
}
else:
if 'user' not in kwargs:
kwargs['user'] = username
if 'password' not in kwargs:
kwargs['password'] = passwd
tdLog.debug(f"connect to {username}:{passwd} with kwargs: {kwargs}")
testconn = taos.connect(**kwargs)
self.cursor = testconn.cursor()
def prepare(self, dbname="db", drop=True, **kwargs):
"""
Prepares the database by optionally dropping it if it exists, creating it, and setting it as the active database.
Args:
dbname (str, optional): The name of the database to be prepared. Defaults to "db".
drop (bool, optional): If True, drops the database if it exists before creating it. Defaults to True.
**kwargs: Additional keyword arguments to be included in the database creation statement. If duration is not provided, it defaults to 100.
Returns:
None
Raises:
None
"""
tdLog.debug(f"prepare database:{dbname}")
s = "reset query cache"
try:
self.cursor.execute(s)
except:
tdLog.notice("'reset query cache' is not supported")
if drop:
s = f"drop database if exists {dbname}"
self.cursor.execute(s)
s = f"create database {dbname}"
for k, v in kwargs.items():
if isinstance(v, str):
s += f" {k} '{v}'"
else:
s += f" {k} {v}"
if "duration" not in kwargs:
s += " duration 100"
if "replica" not in kwargs:
s += f" replica {self.replica}"
tdLog.debug(f"create database cmd: {s}")
self.cursor.execute(s)
s = f"use {dbname}"
self.cursor.execute(s)
time.sleep(2)
def queryAndCheckResult(self, sql_list, expect_result_list, dbPrecision=""):
"""
Executes a list of SQL queries and checks the results against the expected results.
Args:
sql_list (list): The list of SQL queries to be executed.
expect_result_list (list): The list of expected results corresponding to each SQL query.
Returns:
None
Raises:
Exception: If the execution of any SQL query fails or if the results do not match the expected results.
"""
try:
for index in range(len(sql_list)):
self.query(sql_list[index])
if len(expect_result_list[index]) == 0:
self.checkRows(0)
else:
self.checkRows(len(expect_result_list[index]))
for row in range(len(expect_result_list[index])):
for col in range(len(expect_result_list[index][row])):
self.checkData(
row, col, expect_result_list[index][row][col], dbPrecision=dbPrecision
)
except Exception as ex:
raise (ex)
def query(
self,
sql,
row_tag=None,
queryTimes=10,
count_expected_res=None,
show=False,
exit=True,
):
"""
Executes a SQL query and fetches the results.
Args:
sql (str): The SQL query to be executed.
row_tag (optional): If provided, the method will return the fetched results. Defaults to None.
queryTimes (int, optional): The number of times to attempt the query in case of failure. Defaults to 10.
count_expected_res (optional): If provided, the method will repeatedly execute the query until the first result matches this value or the queryTimes limit is reached. Defaults to None.
show (bool, optional): If True, the SQL statement will be logged before execution. Defaults to False.
Returns:
int: The number of rows fetched if row_tag is not provided.
list: The fetched results if row_tag is provided.
Raises:
Exception: If the query fails after the specified number of attempts.
"""
if show:
tdLog.info(sql)
self.sql = sql
i = 1
while i <= queryTimes:
try:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
self.queryRows = len(self.queryResult)
self.queryCols = len(self.cursor.description)
if count_expected_res is not None:
counter = 0
while count_expected_res != self.queryResult[0][0]:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
if counter < queryTimes:
counter += 0.5
time.sleep(0.5)
else:
return False
if row_tag:
return self.queryResult
return self.queryRows
except Exception as e:
if i == queryTimes:
if exit:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.error("%s(%d) failed: sql:%s, %s" % args)
raise
else:
return False
if exit:
tdLog.notice("Try to query again, query times: %d " % i)
i += 1
time.sleep(1)
pass
def querySuccessailed(
self,
sql,
row_tag=None,
queryTimes=10,
count_expected_res=None,
expectErrInfo=None,
fullMatched=True,
):
self.sql = sql
i = 1
while i <= queryTimes:
try:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
self.queryRows = len(self.queryResult)
self.queryCols = len(self.cursor.description)
if count_expected_res is not None:
counter = 0
while count_expected_res != self.queryResult[0][0]:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
if counter < queryTimes:
counter += 0.5
time.sleep(0.5)
else:
return False
tdLog.info("query is success")
time.sleep(1)
continue
except Exception as e:
tdLog.notice("Try to query again, query times: %d " % i)
filename, lineno = _fast_caller(1)
if i < queryTimes:
error_info = repr(e)
print(error_info)
self.error_info = ",".join(
error_info[error_info.index("(") + 1 : -1].split(",")[:-1]
).replace("'", "")
self.queryRows = 0
self.queryCols = 0
self.queryResult = None
if fullMatched:
if expectErrInfo != None:
if expectErrInfo == self.error_info:
tdLog.info(
"sql:%s, expected expectErrInfo '%s' occured"
% (sql, expectErrInfo)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, expectErrInfo '%s' occured, but not expected expectErrInfo '%s'"
% (
filename,
lineno,
sql,
self.error_info,
expectErrInfo,
)
)
else:
if expectErrInfo != None:
if expectErrInfo in self.error_info:
tdLog.info(
"sql:%s, expected expectErrInfo '%s' occured"
% (sql, expectErrInfo)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, expectErrInfo %s occured, but not expected expectErrInfo '%s'"
% (
filename,
lineno,
sql,
self.error_info,
expectErrInfo,
)
)
return self.error_info
elif i == queryTimes:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
i += 1
time.sleep(1)
pass
def executeTimes(self, sql, times):
"""
Executes a SQL statement a specified number of times.(Not used)
Args:
sql (str): The SQL statement to be executed.
times (int): The number of times to execute the SQL statement.
Returns:
int: The number of affected rows from the last execution.
Raises:
None
"""
for i in range(times):
try:
return self.cursor.execute(sql)
except BaseException:
time.sleep(1)
continue
def execute(self, sql, queryTimes=10, show=False):
"""
Executes a SQL statement.
Args:
sql (str): The SQL statement to be executed.
queryTimes (int, optional): The number of times to attempt the execution in case of failure. Defaults to 10.
show (bool, optional): If True, the SQL statement will be logged before execution. Defaults to False.
Returns:
int: The number of affected rows.
Raises:
Exception: If the execution fails after the specified number of attempts.
"""
self.sql = sql
if show:
tdLog.info(sql)
i = 1
while i <= queryTimes:
try:
self.affectedRows = self.cursor.execute(sql)
return self.affectedRows
except Exception as e:
tdLog.notice("Try to execute sql again, execute times: %d " % i)
if i == queryTimes:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
i += 1
time.sleep(1)
pass
# execute many sql
def executes(self, sqls, queryTimes=30, show=False):
"""
Executes a list of SQL statements.
Args:
sqls (list): The list of SQL statements to be executed.
queryTimes (int, optional): The number of times to attempt the execution in case of failure. Defaults to 30.
show (bool, optional): If True, each SQL statement will be logged before execution. Defaults to False.
Returns:
None
Raises:
Exception: If the execution of any SQL statement fails after the specified number of attempts.
"""
for sql in sqls:
self.execute(sql, queryTimes, show)
def waitedQuery(self, sql, expectedRows, timeout):
"""
Executes a SQL query and waits until the expected number of rows is retrieved or the timeout is reached.
Args:
sql (str): The SQL query to be executed.
expectedRows (int): The expected number of rows to be retrieved.
timeout (int): The maximum time to wait (in seconds) for the expected number of rows to be retrieved.
Returns:
tuple: A tuple containing the number of rows retrieved and the time taken (in seconds).
Raises:
Exception: If the query execution fails.
"""
tdLog.info(
"sql: %s, try to retrieve %d rows in %d seconds"
% (sql, expectedRows, timeout)
)
self.sql = sql
try:
for i in range(timeout):
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
self.queryRows = len(self.queryResult)
self.queryCols = len(self.cursor.description)
tdLog.info(
"sql: %s, try to retrieve %d rows,get %d rows"
% (sql, expectedRows, self.queryRows)
)
if self.queryRows >= expectedRows:
return (self.queryRows, i)
time.sleep(1)
except Exception as e:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
return (self.queryRows, timeout)
def query_success_failed(self, sql, row_tag=None, queryTimes=10, count_expected_res=None, expectErrInfo = None, fullMatched = True):
"""Executes a SQL query with retry mechanism and handles both successful and failed scenarios.
This method attempts to execute a SQL query multiple times, handling both successful
executions and expected error conditions. It's particularly useful for testing
scenarios where queries might initially fail but eventually succeed, or for
validating specific error conditions.
Args:
sql (str): The SQL query statement to be executed.
row_tag (optional): If provided, the method will return the fetched results
instead of just the row count. Defaults to None.
queryTimes (int, optional): Maximum number of retry attempts if the query fails.
Defaults to 10.
count_expected_res (optional): If provided, the method will repeatedly execute
the query until the first result matches this value
or retry limit is reached. Defaults to None.
expectErrInfo (str, optional): Expected error message to validate against when
query fails. If None, any error is acceptable.
Defaults to None.
fullMatched (bool, optional): If True, performs exact string matching for error
messages. If False, performs partial string matching
(contains). Defaults to True.
Returns:
str: Error information string if an expected error occurs and query fails
within retry attempts.
None: If query succeeds or if unexpected error occurs and reaches retry limit.
Raises:
Exception: If query fails after all retry attempts and the error is not expected
or doesn't match the expected error pattern.
"""
self.sql = sql
i=1
while i <= queryTimes:
try:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
self.queryRows = len(self.queryResult)
self.queryCols = len(self.cursor.description)
if count_expected_res is not None:
counter = 0
while count_expected_res != self.queryResult[0][0]:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
if counter < queryTimes:
counter += 0.5
time.sleep(0.5)
else:
return False
tdLog.info("query is success")
time.sleep(1)
continue
except Exception as e:
tdLog.notice("Try to query again, query times: %d "%i)
if i < queryTimes:
error_info = repr(e)
print(error_info)
self.error_info = ','.join(error_info[error_info.index('(')+1:-1].split(",")[:-1]).replace("'","")
self.queryRows = 0
self.queryCols = 0
self.queryResult = None
if fullMatched:
if expectErrInfo != None:
if expectErrInfo == self.error_info:
tdLog.info("sql:%s, expected expectErrInfo '%s' occured" % (sql, expectErrInfo))
else:
filename, lineno = _fast_caller(1)
tdLog.exit("%s(%d) failed: sql:%s, expectErrInfo '%s' occured, but not expected expectErrInfo '%s'" % (filename, lineno, sql, self.error_info, expectErrInfo))
else:
if expectErrInfo != None:
if expectErrInfo in self.error_info:
tdLog.info("sql:%s, expected expectErrInfo '%s' occured" % (sql, expectErrInfo))
else:
filename, lineno = _fast_caller(1)
tdLog.exit("%s(%d) failed: sql:%s, expectErrInfo %s occured, but not expected expectErrInfo '%s'" % (filename, lineno, sql, self.error_info, expectErrInfo))
return self.error_info
elif i == queryTimes:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
i+=1
time.sleep(1)
pass
def isErrorSql(self, sql):
"""
Executes a SQL statement and checks if it results in an error.(Not used)
Args:
sql (str): The SQL statement to be executed.
Returns:
bool: True if the SQL statement results in an error, False otherwise.
Raises:
None
"""
err_flag = True
try:
self.cursor.execute(sql)
except BaseException:
err_flag = False
return False if err_flag else True
def errors(
self, sql_list, expected_error_id_list=None, expected_error_info_list=None
):
"""
Executes a list of SQL queries and checks for expected errors.
Args:
sql_list (list): The list of SQL queries to be executed.
expected_error_id_list (list, optional): The list of expected error numbers corresponding to each SQL query. Defaults to None.
expected_error_info_list (list, optional): The list of expected error information corresponding to each SQL query. Defaults to None.
Returns:
None
Raises:
SystemExit: If the SQL list is empty, if the execution of any SQL query fails, if the expected error does not occur, or if the error information does not match the expected information.
"""
try:
if len(sql_list) > 0:
for i in range(len(sql_list)):
if expected_error_id_list and expected_error_info_list:
self.error(
sql_list[i],
expected_error_id_list[i],
expected_error_info_list[i],
)
elif expected_error_id_list:
self.error(sql_list[i], expectedErrno=expected_error_id_list[i])
elif expected_error_info_list:
self.error(
sql_list[i], expectErrInfo=expected_error_info_list[i]
)
else:
self.error(sql_list[i])
else:
tdLog.exit("sql list is empty")
except Exception as ex:
tdLog.exit("Failed to execute sql list: %s, error: %s" % (sql_list, ex))
def is_err_sql(self, sql):
"""Checks if a SQL statement will result in an error when executed.
This method executes the provided SQL statement and determines whether it
causes an exception. It's useful for testing error conditions and validating
that certain SQL statements should fail.
Args:
sql (str): The SQL statement to be tested for errors.
Returns:
bool: False if the SQL statement executes successfully without errors,
True if the SQL statement results in an error/exception.
Raises:
None: This method catches all exceptions internally and returns a boolean
result instead of raising exceptions.
"""
err_flag = True
try:
self.cursor.execute(sql)
except BaseException:
err_flag = False
return False if err_flag else True
def no_error(self, sql):
"""_summary_
Args:
sql (_type_): _description_
"""
expectErrOccurred = False
try:
self.cursor.execute(sql)
except BaseException as e:
expectErrOccurred = True
self.errno = e.errno
error_info = repr(e)
self.error_info = ','.join(error_info[error_info.index('(') + 1:-1].split(",")[:-1]).replace("'", "")
if expectErrOccurred:
filename, lineno = _fast_caller(1)
tdLog.exit("%s(%d) failed: sql:%s, unexpect error '%s' occurred" % (filename, lineno, sql, self.error_info))
else:
tdLog.info("sql:%s, check passed, no ErrInfo occurred" % (sql))
def error(
self, sql, expectedErrno=None, expectErrInfo=None, fullMatched=True, show=False
):
"""
Executes a SQL statement and checks for expected errors.
Args:
sql (str): The SQL statement to be executed.
expectedErrno (int, optional): The expected error number. Defaults to None.
expectErrInfo (str, optional): The expected error information. Defaults to None.
fullMatched (bool, optional): If True, checks for exact matches of the expected error information. Defaults to True.
show (bool, optional): If True, the SQL statement will be logged before execution. Defaults to False.
Returns:
str: The error information if an error occurs.
Raises:
SystemExit: If the expected error does not occur or if the error information does not match the expected information.
"""
filename, lineno = _fast_caller(1)
expectErrNotOccured = True
if show:
tdLog.info(sql)
try:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
self.queryRows = len(self.queryResult)
self.queryCols = len(self.cursor.description)
except BaseException as e:
expectErrNotOccured = False
self.errno = e.errno
error_info = repr(e)
self.error_info = ",".join(
error_info[error_info.index("(") + 1 : -1].split(",")[:-1]
).replace("'", "")
# self.error_info = (','.join(error_info.split(",")[:-1]).split("(",1)[1:][0]).replace("'","")
if expectErrNotOccured:
tdLog.exit(
"%s(%d) failed: sql:%s, expect error not occured"
% (filename, lineno, sql)
)
else:
filename, lineno = _fast_caller(1)
self.queryRows = 0
self.queryCols = 0
self.queryResult = None
if fullMatched:
if expectedErrno != None:
expectedErrno_rest = expectedErrno & 0x0000FFFF
if expectedErrno == self.errno or expectedErrno_rest == (self.errno & 0x0000FFFF):
tdLog.info(
"sql:%s, expected errno %s occured" % (sql, expectedErrno)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, errno '%s' occured, but not expected errno '%s'"
% (
filename,
lineno,
sql,
self.errno,
expectedErrno,
)
)
if expectErrInfo != None:
if expectErrInfo == self.error_info:
tdLog.info(
"sql:%s, expected ErrInfo '%s' occured"
% (sql, expectErrInfo)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, ErrInfo '%s' occured, but not expected ErrInfo '%s'"
% (
filename,
lineno,
sql,
self.error_info,
expectErrInfo,
)
)
else:
if expectedErrno != None:
expectedErrno_rest = expectedErrno & 0x0000FFFF
if expectedErrno in self.errno or expectedErrno_rest in self.errno:
tdLog.info(
"sql:%s, expected errno %s occured" % (sql, expectedErrno)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, errno '%s' occured, but not expected errno '%s'"
% (
filename,
lineno,
sql,
self.errno,
expectedErrno,
)
)
if expectErrInfo != None:
if expectErrInfo in self.error_info:
tdLog.info(
"sql:%s, expected ErrInfo '%s' occured"
% (sql, expectErrInfo)
)
else:
tdLog.exit(
"%s(%d) failed: sql:%s, ErrInfo %s occured, but not expected ErrInfo '%s'"
% (
filename,
lineno,
sql,
self.error_info,
expectErrInfo,
)
)
return self.error_info
# loop call error
def waitError(self, sql, loop=3, sleepMs=1000, expectedErrno=None, expectErrInfo=None, fullMatched=True, show=False):
for i in range(loop):
self.error(sql, expectedErrno, expectErrInfo, fullMatched, show)
time.sleep(sleepMs / 1000)
def noError(self, sql):
unexpectErrOccurred = False
try:
self.cursor.execute(sql)
except BaseException as e:
unexpectErrOccurred = True
self.errno = e.errno
error_info = repr(e)
self.error_info = ",".join(
error_info[error_info.index("(") + 1 : -1].split(",")[:-1]
).replace("'", "")
if unexpectErrOccurred:
filename, lineno = _fast_caller(1)
tdLog.exit(
"%s(%d) failed: sql:%s, unexpect error '%s' occurred"
% (filename, lineno, sql, self.error_info)
)
else:
tdLog.info("sql:%s, check passed, no ErrInfo occurred" % (sql))
def getData(self, row, col):
"""
Retrieves the data at the specified row and column from the last query result.
Args:
row (int): The row index of the data to be retrieved.
col (int): The column index of the data to be retrieved.
Returns:
The data at the specified row and column.
Raises:
SystemExit: If the specified row or column is out of range.
"""
self.checkRowCol(row, col)
return self.queryResult[row][col]
def getDataWithOutCheck(self, row, col):
"""
Retrieves the data at the specified row and column from the last query result.
Args:
row (int): The row index of the data to be retrieved.
col (int): The column index of the data to be retrieved.
Returns:
The data at the specified row and column.
Raises:
IndexError: If the specified row or column is out of range.
"""
return self.queryResult[row][col]
def getColData(self, col):
"""
Retrieves all data from the specified column in the last query result.
Args:
col (int): The column index of the data to be retrieved.
Returns:
list: A list containing all data from the specified column.
Raises:
None
"""
colDatas = []
for i in range(self.queryRows):
colDatas.append(self.queryResult[i][col])
return colDatas
def getRowData(self, row):
"""
Retrieves all data from the specified row in the last query result.
Args:
row (int): The row index of the data to be retrieved.
Returns:
list: A list containing all data from the specified row.
Raises:
SystemExit: If the specified row is out of range.
"""
if row >= self.queryRows:
return None
return self.queryResult[row]
def getResult(self, sql, exit=True):
"""
Executes a SQL query and fetches the results.
Args:
sql (str): The SQL query to be executed.
Returns:
list: The fetched results.
Raises:
Exception: If the query execution fails.
"""
self.sql = sql
try:
self.cursor.execute(sql)
self.queryResult = self.cursor.fetchall()
except Exception as e:
if exit:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
else:
return []
return self.queryResult
def getVariable(self, search_attr):
"""
Retrieves the value of a specified variable from the database.
Args:
search_attr (str): The name of the variable to be retrieved.
Returns:
tuple: A tuple containing the value of the specified variable and the list of all variables.
Raises:
Exception: If the query execution fails.
"""
try:
sql = "show variables"
param_list = self.query(sql, row_tag=True)
for param in param_list:
if param[0] == search_attr:
return param[1], param_list
except Exception as e:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
def getColNameList(self, sql, col_tag=None):
"""
Executes a SQL query and retrieves the column names and optionally the column types.
Args:
sql (str): The SQL query to be executed.
col_tag (optional): If provided, the method will return both column names and column types. Defaults to None.
Returns:
list: A list containing the column names.
tuple: A tuple containing two lists - the column names and the column types, if col_tag is provided.
Raises:
Exception: If the query execution fails.
"""
self.sql = sql
try:
col_name_list = []
col_type_list = []
self.cursor.execute(sql)
for query_col in self.cursor.description:
col_name_list.append(query_col[0])
col_type_list.append(query_col[1])
except Exception as e:
filename, lineno = _fast_caller(1)
args = (filename, lineno, sql, repr(e))
tdLog.notice("%s(%d) failed: sql:%s, %s" % args)
raise Exception(repr(e))
if col_tag:
return col_name_list, col_type_list
return col_name_list
def getRows(self):
"""
Retrieves the number of rows fetched by the last query.
Args:
None
Returns:
int: The number of rows fetched by the last query.
Raises:
None
"""
return self.queryRows
def getCols(self):
"""
Retrieves the number of cols fetched by the last query.
Args:
None
Returns:
int: The number of cols fetched by the last query.
Raises:
None
"""
return self.queryCols
# get first value
def getFirstValue(self, sql):
"""
Executes a SQL query and retrieves the first value in the result.
Args:
sql (str): The SQL query to be executed.
Returns:
The first value in the result.
Raises:
Exception: If the query execution fails.
"""
self.query(sql)
return self.getData(0, 0)
#
# check session
#
def checkRows(self, expectedRows, show=False):
"""
Checks if the number of rows fetched by the last query matches the expected number of rows.
Args:
expectedRows (int): The expected number of rows.
Returns:
bool: True if the number of rows matches the expected number, otherwise it exits the program.
Raises:
SystemExit: If the number of rows does not match the expected number.
"""
return self.checkEqual(self.queryRows, expectedRows, show=show)
def checkRowsV2(
self,
expectedRows: int,
operator: Literal["<", "<=", ">", ">=", "==", "!="] = "==",
show: bool = True,
) -> bool:
"""
Verify if the number of rows returned by SQL query meets the expected condition.
Args:
expectedRows : int
The expected number of rows to compare against
operator : str, optional
Comparison operator ('<', '<=', '>', '>=', '==', '!='),
defaults to '<'
show : bool, optional
Whether to print the verification result, defaults to True
Returns:
bool
True if the actual row count meets the expected condition,
False otherwise
Raises:
ValueError: If invalid operator is provided
Usage:
assert checker.checkRows(15, operator="<") # Verify if less than 15 rows
"""
actualRows = self.queryRows
try:
# Perform comparison based on operator
if operator == "<":
result = actualRows < expectedRows
elif operator == "<=":
result = actualRows <= expectedRows
elif operator == ">":
result = actualRows > expectedRows
elif operator == ">=":
result = actualRows >= expectedRows
elif operator == "==":
result = actualRows == expectedRows
elif operator == "!=":
result = actualRows != expectedRows
else:
raise ValueError(f"Unsupported comparison operator: {operator}")
if show:
tdLog.info(
f"Actual rows: {actualRows}, expected rows: {expectedRows}, comparison: {operator}, result: {'PASS' if result else 'FAIL'}"
)
return result
except Exception as e:
tdLog.error(f"checkRows failed: {str(e)}")
return False
def checkRowsNotExited(self, expectedRows):
"""
Check if the query rows is equal to the expected rows
Args:
expectedRows: The expected number of rows.
Returns:
bool: Returns True if the actual number of rows matches the expected number, otherwise returns False.
"""
if self.queryRows == expectedRows:
return True
else:
return False
def checkRowsRange(self, excepte_row_list):
"""
Checks if the number of rows fetched by the last query is within the expected range.(Not used)
Args:
excepte_row_list (list): A list of expected row counts.
Returns:
bool: True if the number of rows is within the expected range, otherwise it exits the program.
Raises:
SystemExit: If the number of rows is not within the expected range.
"""
if self.queryRows in excepte_row_list:
tdLog.info(
f"sql:{self.sql}, queryRows:{self.queryRows} in expect:{excepte_row_list}"
)
return True
else:
filename, lineno = _fast_caller(1)
tdLog.exit(
f"{filename}({lineno}) failed: sql:{self.sql}, queryRows:{self.queryRows} not in expect:{excepte_row_list}"
)
def checkCols(self, expectCols):
"""
Checks if the number of columns fetched by the last query matches the expected number of columns.
Args:
expectCols (int): The expected number of columns.
Returns:
None
Raises:
SystemExit: If the number of columns does not match the expected number.
"""
if self.queryCols == expectCols:
tdLog.info(
"sql:%s, queryCols:%d == expect:%d"
% (self.sql, self.queryCols, expectCols)
)
else:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
self.queryCols,
expectCols,
)
tdLog.exit("%s(%d) failed: sql:%s, queryCols:%d != expect:%d" % args)
def checkRowCol(self, row, col, exit=True):
"""
Checks if the specified row and column indices are within the range of the last query result.
Args:
row (int): The row index to be checked.
col (int): The column index to be checked.
Returns:
None
Raises:
SystemExit: If the specified row or column index is out of range.
"""
if row < 0:
filename, lineno = _fast_caller(2)
args = (filename, lineno, self.sql, row)
if exit:
tdLog.exit("%s(%d) failed: sql:%s, row:%d is smaller than zero" % args)
else:
return False
if col < 0:
filename, lineno = _fast_caller(2)
args = (filename, lineno, self.sql, row)
if exit:
tdLog.exit("%s(%d) failed: sql:%s, col:%d is smaller than zero" % args)
else:
return False
if row > self.queryRows:
filename, lineno = _fast_caller(2)
args = (filename, lineno, self.sql, row, self.queryRows)
if exit:
tdLog.exit("%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args)
else:
return False
if col > self.queryCols:
filename, lineno = _fast_caller(2)
args = (filename, lineno, self.sql, col, self.queryCols)
if exit:
tdLog.exit("%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args)
else:
return False
return True
def checkDataType(self, row, col, dataType):
"""
Checks if the data type at the specified row and column matches the expected data type.
Args:
row (int): The row index of the data to be checked.
col (int): The column index of the data to be checked.
dataType (str): The expected data type.
Returns:
bool: True if the data type matches the expected data type, otherwise False.
Raises:
SystemExit: If the specified row or column index is out of range.
"""
self.checkRowCol(row, col)
return self.cursor.istype(col, dataType)
def checkFloatString(self, row, col, data, show=False):
if row >= self.queryRows:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, row + 1, self.queryRows)
tdLog.exit(
"%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args
)
if col >= self.queryCols:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, col + 1, self.queryCols)
tdLog.exit(
"%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args
)
self.checkRowCol(row, col)
val = float(self.queryResult[row][col])
if abs(data) >= 1 and abs((val - data) / data) <= 0.000001:
if show:
tdLog.info("check successfully")
elif abs(data) < 1 and abs(val - data) <= 0.000001:
if show:
tdLog.info("check successfully")
else:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s" % args
)
def compareData(self, row, col, data, show=False):
return self.checkData(row, col, data, show, False)
def checkData(self, row, col, data, show=False, exit=True, dbPrecision="", tolerance: float = 0.0):
"""
Checks if the data at the specified row and column matches the expected data.
Args:
row (int): The row index of the data to be checked.
col (int): The column index of the data to be checked.
data: The expected data to be compared with.
show (bool, optional): If True, logs a message when the check is successful. Defaults to False.
Returns:
None
Raises:
SystemExit: If the data at the specified row and column does not match the expected data.
"""
if row >= self.queryRows:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row + 1,
self.queryRows,
)
tdLog.exit(
"%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args
)
else:
return False
if col >= self.queryCols:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
col + 1,
self.queryCols,
)
tdLog.exit(
"%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args
)
else:
return False
if not self.checkRowCol(row, col, exit):
return False
if self.queryResult[row][col] != data:
if self.cursor.istype(col, "TIMESTAMP"):
# suppose user want to check nanosecond timestamp if a longer data passed``
if isinstance(data, str):
if len(data) >= 28:
if self.queryResult[row][col] == _parse_ns_timestamp(data):
if show:
tdLog.info("check successfully")
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
else:
# print(f"{self.queryResult[row][col]}")
real = self.queryResult[row][col]
if real is None:
# none
if str(real) == data:
if show:
tdLog.info("check successfully")
elif isinstance(real, datetime.datetime):
if platform.system().lower() == "windows":
dt_expected = _parse_datetime(data)
# 补齐 tzinfo，避免 Windows astimezone 报错
if real.tzinfo is None:
real = real.replace(tzinfo=datetime.timezone.utc)
if dt_expected.tzinfo is None:
dt_expected = dt_expected.replace(tzinfo=datetime.timezone.utc)
if real.astimezone(datetime.timezone.utc) == dt_expected.astimezone(datetime.timezone.utc):
if show:
tdLog.info("check successfully")
else:
if real.astimezone(datetime.timezone.utc) == _parse_datetime(
data
).astimezone(datetime.timezone.utc):
if show:
tdLog.info("check successfully")
else:
if exit:
caller = inspect.getframeinfo(
inspect.stack()[1][0]
)
args = (
caller.filename,
caller.lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
return True
elif isinstance(data, int):
if(dbPrecision == ""):
if len(str(data)) == 16:
precision = "us"
elif len(str(data)) == 13:
precision = "ms"
elif len(str(data)) == 19:
precision = "ns"
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
f"%s(%d) failed: sql:%s row:%d col:%d len(str(data)):{len(str(data))} data:%s != expect:%s"
% args
)
else:
return False
else:
precision = dbPrecision
success = False
if precision == "ms":
dt_obj = self.queryResult[row][col]
ts = int(
int(
(
dt_obj
- datetime.datetime.fromtimestamp(0, dt_obj.tzinfo)
).total_seconds()
)
* 1000
) + int(dt_obj.microsecond / 1000)
if ts == data:
success = True
elif precision == "us":
dt_obj = self.queryResult[row][col]
ts = int(
int(
(
dt_obj
- datetime.datetime.fromtimestamp(0, dt_obj.tzinfo)
).total_seconds()
)
* 1e6
) + int(dt_obj.microsecond)
if ts == data:
success = True
elif precision == "ns":
if data == self.queryResult[row][col]:
success = True
if success:
if show:
tdLog.info("check successfully")
return True
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
elif isinstance(data, datetime.datetime):
dt_obj = self.queryResult[row][col]
delt_data = data - datetime.datetime.fromtimestamp(0, data.tzinfo)
delt_result = self.queryResult[row][
col
] - datetime.datetime.fromtimestamp(
0, self.queryResult[row][col].tzinfo
)
if delt_data == delt_result:
if show:
tdLog.info("check successfully")
return True
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
return False
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
if str(self.queryResult[row][col]) == str(data):
# tdLog.info(f"sql:{self.sql}, row:{row} col:{col} data:{self.queryResult[row][col]} == expect:{data}")
if show:
tdLog.info("check successfully")
return True
elif isinstance(data, float):
if (
abs(data) >= 1
and abs((self.queryResult[row][col] - data) / data) <= 0.000001
):
# tdLog.info(f"sql:{self.sql}, row:{row} col:{col} data:{self.queryResult[row][col]} == expect:{data}")
if show:
tdLog.info("check successfully")
elif (
abs(data) < 1 and abs(self.queryResult[row][col] - data) <= 0.000001
):
# tdLog.info(f"sql:{self.sql}, row:{row} col:{col} data:{self.queryResult[row][col]} == expect:{data}")
if show:
tdLog.info("check successfully")
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
return True
elif isinstance(data, int) and tolerance != 0.0:
bound1 = data
bound2 = data * (1 + tolerance)
lower = min(bound1, bound2)
upper = max(bound1, bound2)
actual = self.queryResult[row][col]
tdLog.info(f"Tolerance check: lower={lower}, actual={actual}, upper={upper}")
if lower <= actual <= upper:
if show:
tdLog.info("check successfully")
return True
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
lower,
upper,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s not in range [%.6f, %.6f]"
% args
)
else:
return False
else:
if exit:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
row,
col,
self.queryResult[row][col],
data,
)
tdLog.exit(
"%s(%d) failed: sql:%s row:%d col:%d data:%s != expect:%s"
% args
)
else:
return False
if show:
tdLog.info("check successfully")
return True
def checkDataV2(self, row, col, data, show=False, operator="==") -> bool:
"""
Compare the data at the specified row and column with the expected data.
Args:
row (int): The row index of the data to be checked.
col (int): The column index of the data to be checked.
data: The expected data to compare against.
show (bool, optional): If True, logs a message when the check is successful. Defaults to False.
operator (str, optional): The operator to use for comparison. Defaults to "==".
Returns:
bool: True if the comparison is successful, False otherwise.
Usage:
assert self.checkDataV2(row, col, data, show=True, operator="==")
assert self.checkDataV2(row, col, data, show=True, operator="<") # means actual value is less than data
"""
if row >= self.queryRows:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, row + 1, self.queryRows)
tdLog.error(
"%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args
)
return False
if col >= self.queryCols:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, col + 1, self.queryCols)
tdLog.error(
"%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args
)
return False
self.checkRowCol(row, col)
try:
# 获取当前单元格的数据
actual_value = self.queryResult[row][col]
# 转换函数：将值统一转换为可比较的类型
def to_timestamp(value) -> Optional[int]:
"""将任意时间格式转换为纳秒级时间戳"""
if value is None:
return None
# 处理数值时间戳
if isinstance(value, (int, float)):
if value > 1e18:
return int(value)
elif value > 1e15:
tdLog.info(f"value={value}")
tdLog.info(f"int(value * 1000)={int(value * 1000)}")
return int(value * 1000)
elif value > 1e12:
return int(value * 1000000)
else:
return int(value * 1000000000)
# 处理字符串时间
elif isinstance(value, str):
if value.isdigit():
if len(value) == 19:
return int(value)
elif len(value) == 16:
return int(value) * 1000
elif len(value) == 13:
return int(value) * 1000000
elif len(value) == 10:
return int(value) * 1000000000
try:
# 解析常规SQL格式（支持纳秒）
pattern = r"(\d{4}-\d{2}-\d{2})\s(\d{2}:\d{2}:\d{2})(\.\d+)?"
time_match = re.match(pattern, value)
if time_match:
date_part, time_part, nano_part = time_match.groups()
nano = int(
nano_part[1:].ljust(9, "0")[:9] if nano_part else 0
)
dt_str = f"{date_part or '1970-01-01'} {time_part}"
dt = datetime.datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
return int(dt.timestamp()) * 1000000000 + nano
except Exception:
pass
raise ValueError(f"无法解析的时间字符串: {value}")
# 处理datetime对象
elif isinstance(value, datetime.datetime):
return (
int(value.timestamp()) * 1000000000 + value.microsecond * 1000
)
raise TypeError(f"不支持的类型: {type(value)}")
# 转换实际值和预期值
if self.cursor.istype(col, "TIMESTAMP"):
converted_actual = to_timestamp(actual_value)
tdLog.debug(f"sql_result={converted_actual}")
converted_expected = to_timestamp(data)
tdLog.debug(f"sql_expect={converted_expected}")
else:
converted_actual = actual_value
converted_expected = data
# 处理None值比较
if converted_actual is None or converted_expected is None:
result = converted_actual == converted_expected
if operator != "==" and operator != "!=":
raise ValueError("None values can only be compared with == or !=")
result = result if operator == "==" else not result
else:
# 根据操作符进行比较
if operator == "<":
result = converted_actual < converted_expected
elif operator == "<=":
result = converted_actual <= converted_expected
elif operator == ">":
result = converted_actual > converted_expected
elif operator == ">=":
result = converted_actual >= converted_expected
elif operator == "==":
result = converted_actual == converted_expected
elif operator == "!=":
result = converted_actual != converted_expected
else:
raise ValueError(f"Unsupported operator: {operator}")
if show:
tdLog.info(
f"Data at row {row}, col {col} actual value={actual_value}, operator={operator}, expected value={data}, result={result}"
)
return result
except Exception as e:
tdLog.error(f"Error comparing data at row {row}, col {col}: {e}")
return False
def checkKeyData(self, key, col, data, show=False):
"""
Checks if the data at the specified key matches the expected data.
Args:
key: The first column to be compared with.
col (int): The column index of the data to be checked.
data: The expected data to be compared with.
show (bool, optional): If True, logs a message when the check is successful. Defaults to False.
Returns:
None
Raises:
SystemExit: If the data of the specified key does not match the expected data.
"""
if col >= self.queryCols:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, col + 1, self.queryCols)
tdLog.exit(
"%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args
)
row = -1
for i in range(self.queryRows):
if self.queryResult[i][col] == data:
row = i
if show:
tdLog.info(f"find key:{key}, row:{row} col:{col}, data:{data}")
if row == -1:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, key, col)
tdLog.exit("%s(%d) failed: sql:%s key:%s col:%d not found" % args)
if show:
tdLog.info("check key successfully")
def checkKeyExist(self, key, show=False):
return self.checkKeyData(key, 0, key, show=show)
def printResult(self, name="results", exit=False, input_result=None, input_sql=""):
if input_result == None:
rows = self.queryRows
cols = self.queryCols
results = self.queryResult
else:
rows = len(input_result)
if rows == 0:
cols = 0
else:
cols = len(input_result[0])
results = input_result
if input_sql == "":
sql = self.sql
else:
sql = input_sql
tdLog.info(f"==== {name}, rows:{rows}, cols:{cols}, sql:{sql}")
for r in range(rows):
data = "==== "
for c in range(cols):
data += f"d[{r}][{c}]={results[r][c]} "
tdLog.info(data)
if exit:
filename, lineno = _fast_caller(1)
tdLog.info(f"{name} {filename}({lineno})")
filename, lineno = _fast_caller(2)
tdLog.exit(f"{name} {filename}({lineno})")
def expectKeyData(self, key, col, data, show=False):
"""
Whether the data at the specified key matches the expected data.
Args:
key: The first column to be compared with.
col (int): The column index of the data to be checked.
data: The expected data to be compared with.
show (bool, optional): If True, logs a message when the check is successful. Defaults to False.
Returns:
Bool
Raises:
None
"""
if col >= self.queryCols:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, col + 1, self.queryCols)
tdLog.info(
"%s(%d) failed: sql:%s, col:%d is larger than queryCols:%d" % args
)
return False
row = -1
for i in range(self.queryRows):
if self.queryResult[i][col] == data:
row = i
tdLog.info(f"find key:{key}, row:{row} col:{col}, data:{data}")
if row == -1:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, key, col)
return False
if show:
tdLog.info("check key successfully")
return True
def checkAssert(self, assertVal, show=False):
"""
Checks if the assertVal is true.
Args:
assertVal: The value to be assert
show (bool, optional): If True, logs a message when the check is successful. Defaults to False.
Returns:
None
Raises:
SystemExit: If the data of the specified key does not match the expected data.
"""
if assertVal != True:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql)
tdLog.exit("%s(%d) failed: sql:%s asserted" % args)
if show:
tdLog.info("check assert successfully")
def checkDataMem(self, sql, mem):
"""
Executes a SQL query and checks if the result matches the expected data.
Args:
sql (str): The SQL query to be executed.
mem (list): The expected data, represented as a list of lists.
Returns:
None
Raises:
SystemExit: If the expected data is not a list of lists, or if the SQL result does not match the expected data.
"""
self.query(sql)
if not isinstance(mem, list):
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql)
tdLog.exit(
"%s(%d) failed: sql:%s, expect data is error, must is array[][]" % args
)
if len(mem) != self.queryRows:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, len(mem), self.queryRows)
tdLog.exit(
"%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args
)
# row, col, data
for row, rowData in enumerate(mem):
for col, colData in enumerate(rowData):
self.checkData(row, col, colData)
tdLog.info("check successfully")
def checkDataMemLoop(self, sql, mem, loopCount=120, waitTime=1):
"""
Loop executes a SQL query and checks if the result matches the expected data.
Args:
sql (str): The SQL query to be executed.
mem (list): The expected data, represented as a list of lists.
Returns:
None
Raises:
SystemExit: If the expected data is not a list of lists, or if the SQL result does not match the expected data.
"""
if not isinstance(mem, list):
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql)
tdLog.exit("%s(%d) failed: sql:%s, expect data is error, must is array[][]" % args)
#
# loop check not exit
#
for i in range(loopCount):
self.query(sql)
if len(mem) == self.queryRows:
passed = True
for row, rowData in enumerate(mem):
for col, colData in enumerate(rowData):
if self.checkData(row, col, colData, exit=False) == False:
passed = False
break
if passed == False:
break
if passed:
# succ return
tdLog.info(f"checkDataMemLoop attempt {i+1}/{loopCount} succeeded")
return
# failed retry next
tdLog.info(f"checkDataMemLoop attempt {i+1}/{loopCount} failed, waiting {waitTime} seconds before retrying...")
time.sleep(waitTime)
#
# last check
#
self.checkDataMem(sql, mem)
def checkDataCsv(self, sql, csvfilePath):
"""
Executes a SQL query and checks if the result matches the expected data from a CSV file.(Not used)
Args:
sql (str): The SQL query to be executed.
csvfilePath (str): The path to the CSV file containing the expected data.
Returns:
None
Raises:
SystemExit: If the CSV file path is invalid, the file is not found, there is an error reading the file,
or if the sql result does not match the expected data from CSV file.
"""
if not isinstance(csvfilePath, str) or len(csvfilePath) == 0:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath)
tdLog.exit("%s(%d) failed: sql:%s, expect csvfile path error:%s" % args)
tdLog.info("read csvfile read begin")
data = []
try:
with open(csvfilePath) as csvfile:
csv_reader = csv.reader(csvfile) # csv.reader read csvfile\
# header = next(csv_reader) # Read the header of each column in the first row
for row in csv_reader: # csv file save to data
data.append(row)
except FileNotFoundError:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath)
tdLog.exit("%s(%d) failed: sql:%s, expect csvfile not find error:%s" % args)
except Exception as e:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath, str(e))
tdLog.exit(
"%s(%d) failed: sql:%s, expect csvfile path:%s, read error:%s" % args
)
tdLog.info("read csvfile read successfully")
self.checkDataMem(sql, data)
def checkDataMemByLine(self, sql, mem):
"""
Executes a SQL query and checks if the result matches the expected data (Same as checkDataMem).
Args:
sql (str): The SQL query to be executed.
mem (list): The expected data, represented as a list of lists.
Returns:
None
Raises:
SystemExit: If the expected data is not a list of lists, or if the SQL result does not match the expected data.
"""
if not isinstance(mem, list):
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql)
tdLog.exit(
"%s(%d) failed: sql:%s, expect data is error, must is array[][]" % args
)
if len(mem) != self.queryRows:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, len(mem), self.queryRows)
tdLog.exit(
"%s(%d) failed: sql:%s, row:%d is larger than queryRows:%d" % args
)
# row, col, data
for row, rowData in enumerate(mem):
for col, colData in enumerate(rowData):
self.checkData(row, col, colData)
tdLog.info("check %s successfully" % sql)
def checkDataCsvByLine(self, sql, csvfilePath):
"""
Executes a SQL query and checks if the result matches the expected data from a CSV file line by line.
Args:
sql (str): The SQL query to be executed.
csvfilePath (str): The path to the CSV file containing the expected data.
Returns:
None
Raises:
SystemExit: If the CSV file path is invalid, the file is not found, there is an error reading the file,
or if the SQL result does not match the expected data from the CSV file.
"""
if not isinstance(csvfilePath, str) or len(csvfilePath) == 0:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath)
tdLog.exit("%s(%d) failed: sql:%s, expect csvfile path error:%s" % args)
self.query(sql)
data = []
tdLog.info("check line %d start" % self.csvLine)
try:
with open(csvfilePath) as csvfile:
skip_rows = self.csvLine
# 计算需要读取的行数
num_rows = self.queryRows
# 读取指定范围的行
df = pd.read_csv(
csvfilePath, skiprows=skip_rows, nrows=num_rows, header=None
)
for index, row in df.iterrows():
data.append(row)
self.csvLine += self.queryRows
except FileNotFoundError:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath)
tdLog.exit("%s(%d) failed: sql:%s, expect csvfile not find error:%s" % args)
except Exception as e:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, csvfilePath, str(e))
tdLog.exit(
"%s(%d) failed: sql:%s, expect csvfile path:%s, read error:%s" % args
)
self.checkDataMemByLine(sql, data)
# return true or false replace exit, no print out
def checkRowColNoExist(self, row, col):
"""
Checks if the specified row and column indices are within the range of the last query result without exiting the program.
Args:
row (int): The row index to be checked.
col (int): The column index to be checked.
Returns:
bool: True if the specified row and column indices are within the range, otherwise False.
Raises:
None
"""
if row < 0:
return False
if col < 0:
return False
if row > self.queryRows:
return False
if col > self.queryCols:
return False
return True
# return true or false replace exit, no print out
def checkDataNoExist(self, row, col, data):
"""
Checks if the data at the specified row and column matches the expected data without exiting the program.
Args:
row (int): The row index of the data to be checked.
col (int): The column index of the data to be checked.
data: The expected data to be compared with.
Returns:
bool: True if the data matches the expected data, otherwise False.
Raises:
None
"""
if self.checkRowColNoExist(row, col) == False:
return False
if self.queryResult[row][col] != data:
if self.cursor.istype(col, "TIMESTAMP"):
# suppose user want to check nanosecond timestamp if a longer data passed
if len(data) >= 28:
if pd.to_datetime(self.queryResult[row][col]) == pd.to_datetime(
data
):
return True
else:
if self.queryResult[row][col] == _parse_datetime(data):
return True
return False
if str(self.queryResult[row][col]) == str(data):
return True
elif isinstance(data, float):
if (
abs(data) >= 1
and abs((self.queryResult[row][col] - data) / data) <= 0.000001
):
return True
elif (
abs(data) < 1 and abs(self.queryResult[row][col] - data) <= 0.000001
):
return True
else:
return False
else:
return False
return True
# loop execute sql then sleep(waitTime) , if checkData ok break loop
def checkDataLoop(self, row, col, data, sql, loopCount=10, waitTime=1):
"""
Executes a SQL query in a loop and checks if the data at the specified row and column matches the expected data.
Args:
row (int): The row index of the data to be checked.
col (int): The column index of the data to be checked.
data: The expected data to be compared with.
sql (str): The SQL query to be executed.
loopCount (int): The number of times to execute the SQL query.
waitTime (int): The time to wait (in seconds) between each execution.
Returns:
None
Raises:
Exception: If the query execution fails.
SystemExit: If the data at the specified row and column does not match the expected data.
"""
# loop check util checkData return true
for i in range(loopCount):
self.query(sql)
if self.checkData(row, col, data, exit=False):
return
#print(f"checkData failed, retrying {i} [{row},{col}] {data} ...")
time.sleep(waitTime)
# last check
self.query(sql)
self.checkData(row, col, data)
def checkRowsLoop(self, expectedRows, sql, loopCount, waitTime):
# loop check util checkData return true
for i in range(loopCount):
self.query(sql)
if self.checkRowsNotExited(expectedRows):
return
else:
time.sleep(waitTime)
continue
# last check
self.query(sql)
self.checkRows(expectedRows)
def checkAffectedRows(self, expectAffectedRows):
"""
Checks if the number of affected rows from the last executed SQL statement matches the expected number of affected rows.
Args:
expectAffectedRows (int): The expected number of affected rows.
Returns:
None
Raises:
SystemExit: If the number of affected rows does not match the expected number.
"""
if self.affectedRows != expectAffectedRows:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
self.affectedRows,
expectAffectedRows,
)
tdLog.exit("%s(%d) failed: sql:%s, affectedRows:%d != expect:%d" % args)
tdLog.info(
"sql:%s, affectedRows:%d == expect:%d"
% (self.sql, self.affectedRows, expectAffectedRows)
)
def checkColNameList(self, col_name_list, expect_col_name_list):
"""
Checks if the column names from the last query match the expected column names.
Args:
col_name_list (list): The list of column names from the last query.
expect_col_name_list (list): The list of expected column names.
Returns:
None
Raises:
SystemExit: If the column names do not match the expected column names.
"""
if col_name_list == expect_col_name_list:
tdLog.info(
"sql:%s, col_name_list:%s == expect_col_name_list:%s"
% (self.sql, col_name_list, expect_col_name_list)
)
else:
filename, lineno = _fast_caller(1)
args = (
filename,
lineno,
self.sql,
col_name_list,
expect_col_name_list,
)
tdLog.exit(
"%s(%d) failed: sql:%s, col_name_list:%s != expect_col_name_list:%s"
% args
)
def checkResColNameList(self, expect_col_name_list):
col_name_list = []
col_type_list = []
for query_col in self.cursor.description:
col_name_list.append(query_col[0])
col_type_list.append(query_col[1])
self.checkColNameList(col_name_list, expect_col_name_list)
def __check_equal(self, elm, expect_elm):
if elm == expect_elm:
return True
if isinstance(elm, datetime.datetime) and isinstance(expect_elm, str):
try:
parsed = datetime.datetime.fromisoformat(expect_elm)
return elm == parsed
except ValueError:
return False
if isinstance(expect_elm, datetime.datetime) and isinstance(elm, str):
try:
parsed = datetime.datetime.fromisoformat(elm)
return expect_elm == parsed
except ValueError:
return False
if type(elm) in (list, tuple) and type(expect_elm) in (list, tuple):
if len(elm) != len(expect_elm):
return False
if len(elm) == 0:
return True
for i in range(len(elm)):
flag = self.__check_equal(elm[i], expect_elm[i])
if not flag:
return False
return True
return False
def print_error_frame_info(self, elm, expect_elm, sql=None):
filename, lineno = _fast_caller(1)
print_sql = self.sql if sql is None else sql
args = (filename, lineno, print_sql, elm, expect_elm)
# tdLog.info("%s(%d) failed: sql:%s, elm:%s != expect_elm:%s" % args)
raise Exception("%s(%d) failed: sql:%s, elm:%s != expect_elm:%s" % args)
def checkEqual(self, elm, expect_elm, show=False):
"""
Checks if the given element is equal to the expected element.
Args:
elm: The element to be checked.
expect_elm: The expected element to be compared with.
Returns:
None
Raises:
Exception: If the element does not match the expected element.
"""
if elm == expect_elm:
if show:
tdLog.info(
"sql:%s, elm:%s == expect_elm:%s" % (self.sql, elm, expect_elm)
)
return True
if self.__check_equal(elm, expect_elm):
if show:
tdLog.info(
"sql:%s, elm:%s == expect_elm:%s" % (self.sql, elm, expect_elm)
)
return True
self.print_error_frame_info(elm, expect_elm)
def checkNotEqual(self, elm, expect_elm):
"""
Checks if the given element is not equal to the expected element.
Args:
elm: The element to be checked.
expect_elm: The expected element to be compared with.
Returns:
None
Raises:
Exception: If the element matches the expected element.
"""
if elm != expect_elm:
tdLog.info("sql:%s, elm:%s != expect_elm:%s" % (self.sql, elm, expect_elm))
else:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, elm, expect_elm)
tdLog.info("%s(%d) failed: sql:%s, elm:%s == expect_elm:%s" % args)
raise Exception
def checkGreater(self, elm, expect_elm):
"""Verifies that the first element is greater than the second element.
This method compares two values and ensures that the first value (elm) is
strictly greater than the second value (expect_elm). It's commonly used for
validating query results, performance metrics, or any numeric comparisons
in test cases.
Args:
elm: The actual value to be compared. Can be any comparable type
(int, float, string, etc.).
expect_elm: The expected threshold value that elm should exceed.
Must be the same comparable type as elm.
Returns:
bool: True if elm > expect_elm, False otherwise.
"""
if elm > expect_elm:
tdLog.info("sql:%s, elm:%s > expect_elm:%s" % (self.sql, elm, expect_elm))
return True
else:
filename, lineno = _fast_caller(1)
args = (filename, lineno, self.sql, elm, expect_elm)
tdLog.info("%s(%d) failed: sql:%s, elm:%s <= expect_elm:%s" % args)
self.print_error_frame_info(elm, expect_elm)
return False
# check like select count(*) ... sql
def checkAgg(self, sql, expectCnt):
"""
Executes an aggregate SQL query and checks if the result matches the expected count.
Args:
sql (str): The aggregate SQL query to be executed.
expectCnt (int): The expected count from the aggregate query.
Returns:
None
Raises:
Exception: If the query execution fails.
SystemExit: If the result of the aggregate query does not match the expected count.
"""
self.query(sql)
self.checkData(0, 0, expectCnt)
tdLog.info(f"{sql} expect {expectCnt} ok.")
# expect first value
def checkFirstValue(self, sql, expect):
"""
Executes a SQL query and checks if the first value in the result matches the expected value.
Args:
sql (str): The SQL query to be executed.
expect: The expected value of the first result.
Returns:
None
Raises:
Exception: If the query execution fails.
SystemExit: If the first value in the result does not match the expected value.
"""
self.query(sql)
self.checkData(0, 0, expect)
# colIdx1 value same with colIdx2
def checkSameColumn(self, c1, c2):
"""
Checks if the values in two specified columns are the same for all rows in the last query result.
Args:
c1 (int): The index of the first column to be checked.
c2 (int): The index of the second column to be checked.
Returns:
None
Raises:
SystemExit: If the values in the specified columns are not the same for any row.
"""
for i in range(self.queryRows):
if self.queryResult[i][c1] != self.queryResult[i][c2]:
tdLog.exit(
f"Not same. row={i} col1={c1} col2={c2}. {self.queryResult[i][c1]}!={self.queryResult[i][c2]}"
)
tdLog.info(
f"check {self.queryRows} rows two column value same. column index [{c1},{c2}]"
)
#
# others session
#
def getTimes(self, time_str, precision="ms"):
"""
Converts a time string to a timestamp based on the specified precision.(Not used)
Args:
time_str (str): The time string to be converted. The string should end with a character indicating the time unit (e.g., 's' for seconds, 'm' for minutes).
precision (str, optional): The precision of the timestamp. Can be "ms" (milliseconds), "us" (microseconds), or "ns" (nanoseconds). Defaults to "ms".
Returns:
int: The timestamp in the specified precision.
Raises:
SystemExit: If the time string does not end with a valid time unit character or if the precision is not valid.
"""
if time_str[-1] not in TAOS_TIME_INIT:
filename, lineno = _fast_caller(1)
tdLog.exit(
f"{filename}({lineno}) failed: {time_str} not a standard taos time init"
)
if precision not in TAOS_PRECISION:
filename, lineno = _fast_caller(1)
tdLog.exit(
f"{filename}({lineno}) failed: {precision} not a standard taos time precision"
)
if time_str[-1] == TAOS_TIME_INIT[0]:
times = int(time_str[:-1]) * TIME_NS
if time_str[-1] == TAOS_TIME_INIT[1]:
times = int(time_str[:-1]) * TIME_US
if time_str[-1] == TAOS_TIME_INIT[2]:
times = int(time_str[:-1]) * TIME_MS
if time_str[-1] == TAOS_TIME_INIT[3]:
times = int(time_str[:-1]) * TIME_S
if time_str[-1] == TAOS_TIME_INIT[4]:
times = int(time_str[:-1]) * TIME_M
if time_str[-1] == TAOS_TIME_INIT[5]:
times = int(time_str[:-1]) * TIME_H
if time_str[-1] == TAOS_TIME_INIT[6]:
times = int(time_str[:-1]) * TIME_D
if time_str[-1] == TAOS_TIME_INIT[7]:
times = int(time_str[:-1]) * TIME_W
if time_str[-1] == TAOS_TIME_INIT[8]:
times = int(time_str[:-1]) * TIME_N
if time_str[-1] == TAOS_TIME_INIT[9]:
times = int(time_str[:-1]) * TIME_Y
if precision == "ms":
return int(times)
elif precision == "us":
return int(times * 1000)
elif precision == "ns":
return int(times * 1000 * 1000)
def getType(self, col):
"""
Retrieves the data type of the specified column in the last query result.(Not used)
Args:
col (int): The column index for which the data type is to be retrieved.
Returns:
str: The data type of the specified column.
Raises:
None
"""
if self.cursor.istype(col, "BOOL"):
return "BOOL"
if self.cursor.istype(col, "INT"):
return "INT"
if self.cursor.istype(col, "BIGINT"):
return "BIGINT"
if self.cursor.istype(col, "TINYINT"):
return "TINYINT"
if self.cursor.istype(col, "SMALLINT"):
return "SMALLINT"
if self.cursor.istype(col, "FLOAT"):
return "FLOAT"
if self.cursor.istype(col, "DOUBLE"):
return "DOUBLE"
if self.cursor.istype(col, "BINARY"):
return "BINARY"
if self.cursor.istype(col, "NCHAR"):
return "NCHAR"
if self.cursor.istype(col, "TIMESTAMP"):
return "TIMESTAMP"
if self.cursor.istype(col, "JSON"):
return "JSON"
if self.cursor.istype(col, "TINYINT UNSIGNED"):
return "TINYINT UNSIGNED"
if self.cursor.istype(col, "SMALLINT UNSIGNED"):
return "SMALLINT UNSIGNED"
if self.cursor.istype(col, "INT UNSIGNED"):
return "INT UNSIGNED"
if self.cursor.istype(col, "BIGINT UNSIGNED"):
return "BIGINT UNSIGNED"
def taosdStatus(self, state):
tdLog.sleep(5)
pstate = 0
for i in range(30):
pstate = 0
pl = psutil.pids()
for pid in pl:
try:
if psutil.Process(pid).name() == "taosd":
print("have already started")
pstate = 1
break
except psutil.NoSuchProcess:
pass
if pstate == state:
break
if state or pstate:
tdLog.sleep(1)
continue
pstate = 0
break
args = (pstate, state)
if pstate == state:
tdLog.info("taosd state is %d == expect:%d" % args)
else:
tdLog.exit("taosd state is %d != expect:%d" % args)
pass
def haveFile(self, dir, state):
if os.path.exists(dir) and os.path.isdir(dir):
if not os.listdir(dir):
if state:
tdLog.exit("dir: %s is empty, expect: not empty" % dir)
else:
tdLog.info("dir: %s is empty, expect: empty" % dir)
else:
if state:
tdLog.info("dir: %s is not empty, expect: not empty" % dir)
else:
tdLog.exit("dir: %s is not empty, expect: empty" % dir)
else:
tdLog.exit("dir: %s doesn't exist" % dir)
def createDir(self, dir):
if os.path.exists(dir):
shutil.rmtree(dir)
tdLog.info("dir: %s is removed" % dir)
os.makedirs(dir, 755)
tdLog.info("dir: %s is created" % dir)
pass
def getDbVgroups(self, db_name: str = "test") -> list:
db_vgroups_list = []
self.query(f"show {db_name}.vgroups")
for result in self.queryResult:
db_vgroups_list.append(result[0])
vgroup_nums = len(db_vgroups_list)
tdLog.debug(f"{db_name} has {vgroup_nums} vgroups :{db_vgroups_list}")
self.query("select * from information_schema.ins_vnodes")
return db_vgroups_list
def getCluseterDnodes(self) -> list:
cluset_dnodes_list = []
self.query("show dnodes")
for result in self.queryResult:
cluset_dnodes_list.append(result[0])
self.clust_dnode_nums = len(cluset_dnodes_list)
tdLog.debug(
f"cluster has {len(cluset_dnodes_list)} dnodes :{cluset_dnodes_list}"
)
return cluset_dnodes_list
def redistribute_one_vgroup(
self,
db_name: str = "test",
replica: int = 1,
vgroup_id: int = 1,
useful_trans_dnodes_list: list = [],
):
# redisutribute vgroup {vgroup_id} dnode {dnode_id}
if replica == 1:
dnode_id = random.choice(useful_trans_dnodes_list)
redistribute_sql = f"redistribute vgroup {vgroup_id} dnode {dnode_id}"
elif replica == 3:
selected_dnodes = random.sample(useful_trans_dnodes_list, replica)
redistribute_sql_parts = [f"dnode {dnode}" for dnode in selected_dnodes]
redistribute_sql = f"redistribute vgroup {vgroup_id} " + " ".join(
redistribute_sql_parts
)
else:
raise ValueError(f"Replica count must be 1 or 3,but got {replica}")
tdLog.debug(f"redistributeSql:{redistribute_sql}")
self.query(redistribute_sql)
tdLog.debug("redistributeSql ok")
def redistributeDbAllVgroups(self, db_name: str = "test", replica: int = 1):
db_vgroups_list = self.getDbVgroups(db_name)
cluset_dnodes_list = self.getCluseterDnodes()
useful_trans_dnodes_list = cluset_dnodes_list.copy()
self.query("select * from information_schema.ins_vnodes")
# result: dnode_id|vgroup_id|db_name|status|role_time|start_time|restored|
results = list(self.queryResult)
for vnode_group_id in db_vgroups_list:
for result in results:
print(
f"result[2] is {result[2]}, db_name is {db_name}, result[1] is {result[1]}, vnode_group_id is {vnode_group_id}"
)
if result[2] == db_name and result[1] == vnode_group_id:
tdLog.debug(
f"dbname: {db_name}, vgroup :{vnode_group_id}, dnode is {result[0]}"
)
print(useful_trans_dnodes_list)
useful_trans_dnodes_list.remove(result[0])
tdLog.debug(
f"vgroup :{vnode_group_id},redis_dnode list:{useful_trans_dnodes_list}"
)
self.redistribute_one_vgroup(
db_name, replica, vnode_group_id, useful_trans_dnodes_list
)
useful_trans_dnodes_list = cluset_dnodes_list.copy()
def pause(self):
"""
Pause the execution of the program and wait for enter key. Used for debugging.
Args:
None
Returns:
None
Raises:
None
"""
if os.name == "nt": # Windows
os.system("pause") # 显示 "按任意键继续..."
else: # Linux/macOS
input("press enter to continue...")
def setConnMode(self, mode=0, value=1):
"""
Set Conn Mode
Args:
mode (int, optional): connect mode options.
Returns:
None
Raises:
None
"""
tdLog.info(f"set connection mode:{mode} value:{value}")
self.cursor._connection.set_mode(mode, value)
def checkResultsByFunc(self, sql, func, delay=0.0, retry=300, show=False):
if delay != 0:
time.sleep(delay)
# show sql
tdLog.info(sql)
if retry <= 0:
retry = 1
for loop in range(retry):
self.clearResult()
if self.query(sql, queryTimes=1, exit=False):
if func():
self.printResult(f"check succeed in {loop} seconds")
return
if loop != retry - 1:
if show:
self.printResult(f"check continue {loop} after sleep 1s ...")
time.sleep(1)
self.printResult(f"check failed for {retry} seconds, sql={sql}", exit=True)
def checkResultsByArray(
self, sql, exp_result, exp_sql="", delay=0.0, retry=300, show=False
):
if delay != 0:
time.sleep(delay)
if retry <= 0:
retry = 1
exp_rows = len(exp_result)
exp_cols = 0
if exp_rows > 0:
exp_cols = len(exp_result[0])
for loop in range(retry):
self.clearResult()
if self.query(sql, queryTimes=1, exit=False) and self.getRows() == exp_rows:
res_result = self.queryResult
all_same = True
for r in range(exp_rows):
for c in range(exp_cols):
if not self.compareData(r, c, exp_result[r][c], show):
all_same = False
break;
if not all_same:
break;
if all_same:
self.printResult(
f"check succeed in {loop} seconds", input_result=res_result
)
return
if loop != retry - 1:
if show:
self.printResult("check continue", input_result=res_result)
time.sleep(1)
self.printResult(f"expect results", input_result=exp_result, input_sql=exp_sql)
self.printResult(
f"check failed for {retry} seconds", input_result=res_result, input_sql=sql
)
self.compareResults(res_result, exp_result, show=True)
filename, lineno = _fast_caller(1)
tdLog.info(f"{filename}({lineno}) check result failed")
filename, lineno = _fast_caller(2)
tdLog.exit(f"{filename}({lineno}) check result failed")
def checkResultsBySql(self, sql, exp_sql, delay=0.0, retry=300, show=False):
# sleep
if delay != 0:
time.sleep(delay)
if retry <= 0:
retry = 1
# loop retry
for loop in range(retry):
# clear
self.clearResult()
# query
result = self.getResult(sql, exit=False)
exp_result = self.getResult(exp_sql, exit=False)
# check result
if self.compareResults(result, exp_result):
# success
return
# sleep and retry
if loop != retry - 1:
if show:
self.printResult(f"check continue {loop} after sleep 1s ...")
time.sleep(1)
print(f"result is: {result}")
print(f"exp_result is: {exp_result}")
raise Exception(f"check failed for {retry} seconds, sql={sql}, exp_sql={exp_sql}")
def checkTableType(
self,
dbname,
columns,
tags=0,
stbname=None,
tbname=None,
typename=None,
delay=0.0,
retry=60,
show=False,
):
if tags == 0:
sql = f"select * from information_schema.ins_tables where db_name='{dbname}' and table_name='{tbname}'"
self.checkResultsByFunc(
sql=sql,
func=lambda: self.getRows() == 1
and self.compareData(0, 0, tbname)
and self.compareData(0, 1, dbname)
and self.compareData(0, 3, columns)
and self.compareData(0, 4, stbname)
and self.compareData(0, 9, typename),
delay=delay,
retry=retry,
show=show,
)
else:
sql = f"select * from information_schema.ins_stables where db_name='{dbname}' and stable_name='{stbname}'"
self.checkResultsByFunc(
sql=sql,
func=lambda: self.getRows() == 1
and self.compareData(0, 0, stbname)
and self.compareData(0, 1, dbname)
and self.compareData(0, 3, columns)
and self.compareData(0, 4, tags),
delay=delay,
retry=retry,
show=show,
)
def checkTableSchema(
self,
dbname,
tbname,
schema,
delay=0.0,
retry=60,
show=False,
):
sql = f"desc {dbname}.{tbname}"
self.checkResultsByFunc(
sql=sql,
func=lambda: self.compareSchema(schema),
delay=delay,
retry=retry,
show=show,
)
def compareSchema(self, schema):
row = len(schema)
for r in range(row):
if (
schema[r][0] != self.queryResult[r][0] # field
or schema[r][1] != self.queryResult[r][1] # type
or schema[r][2] != self.queryResult[r][2] # length
or schema[r][3] != self.queryResult[r][3] # note
):
tdLog.info(f"exp_schema[{r}]={schema[r]}, res_schema[{r}]={self.queryResult[r]}")
return False
return True
def compareResults(self, res_result, exp_result, show=False):
exp_rows = len(exp_result)
res_rows = len(res_result)
if exp_rows != res_rows:
if show:
tdLog.info(f"exp_rows={exp_rows}, res_rows={res_rows}")
return False
if exp_rows == 0:
return True
exp_cols = len(exp_result[0])
res_cols = len(res_result[0])
if exp_cols != res_cols:
if show:
tdLog.info(f"exp_cols={exp_cols}, res_cols={res_cols}")
return False
if exp_cols == 0:
return True
for r in range(res_rows):
for c in range(res_cols):
if res_result[r][c] != exp_result[r][c]:
if show:
tdLog.info(
f"res_result[{r}][{c}]={res_result[r][c]}, exp_result[{r}][{c}]={exp_result[r][c]}"
)
return False
return True
def clearResult(self):
self.queryCols = 0
self.queryRows = 0
self.queryResult = []
# insert table with fixed values, return next write ts
def insertFixedVal(self, table, startTs, step, count, cols, fixedVals):
# init
ts = startTs
# loop count
for i in range(count):
sql = f"INSERT INTO {table}({cols}) VALUES({ts},{fixedVals})"
self.execute(sql, show=True)
# next
ts += step
return ts
# gen insert table with fixed values, return next write ts
def genInsertVal(self, startTs, step, count, fixedVals):
# init
ts = startTs
sql = ""
# loop count
for i in range(count):
sql += f" ({ts},{fixedVals})"
# next
ts += step
return ts, sql
# insert now
def insertNow(self, table, sleepS, count, cols, fixedVals):
# loop count
for i in range(count):
sql = f"INSERT INTO {table}({cols}) VALUES(now,{fixedVals})"
self.execute(sql, show=True)
# next
time.sleep(sleepS)
# insert table with order values, only support number cols, return next write ts
def insertOrderVal(self, table, startTs, step, count, cols, orderVals, colStep = 1):
# init
ts = startTs
colsVal = orderVals
# loop count
for i in range(count):
# insert sql
sql = f"INSERT INTO {table}({cols}) VALUES({ts}, {','.join(map(str, colsVal))})"
self.execute(sql, show=True)
# next
ts += step
for j in range(len(colsVal)):
colsVal[j] += colStep
return ts
# flush db
def flushDb(self, dbName):
self.execute(f"flush database {dbName}", show=True)
# delete rows
def deleteRows(self, table, where=None):
"""
Deletes rows from the specified table based on the given condition.
Args:
table (str): The name of the table from which rows are to be deleted.
where (str, optional): The condition for deleting rows. Defaults to None.
Returns:
None
Raises:
SystemExit: If the delete operation fails.
"""
sql = f"DELETE FROM {table}"
if where:
sql += f" WHERE {where}"
self.execute(sql, show=True)
def fieldIndex(self, fieldName):
"""
Retrieves the index of the specified field name in the last query result.
Args:
fieldName (str): The name of the field whose index is to be retrieved.
Returns:
int: The index of the specified field name.
Raises:
SystemExit: If the field name does not exist in the last query result.
"""
for idx, col in enumerate(self.cursor.description):
if col[0].lower() == fieldName.lower():
return idx
filename, lineno = _fast_caller(1)
tdLog.exit(f"{filename}({lineno}) failed: field name {fieldName} not exist")
# global
tdSql = TDSql()
Reference in a new issue View git blame Copy permalink