pytest/lib/common.py - doris - Git at Google

 #!/bin/env python
 # -*- coding: utf-8 -*-
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.

 ###########################################################################
 #   @file common.py
 #   @date 2020-07-14 19:11:39
 #   @brief common function
 ############################################################################
 """
 common test tool
 """
 import re
 import time

 import palo_client
 import palo_config
 import palo_logger
 import palo_job
 import palo_types
 import util

 config = palo_config.config
 LOG = palo_logger.Logger.getLogger()
 L = palo_logger.StructedLogMessage


 def get_client(host=None):
     """get a new client"""
     if host is None:
         host = config.fe_host
     client = palo_client.get_client(host, config.fe_query_port, user=config.fe_user,
                                     password=config.fe_password, http_port=config.fe_http_port)
     return client


 def create_workspace(database_name):
     """get new palo client connect, and create db, and use and return client"""
     client = palo_client.get_client(config.fe_host, config.fe_query_port, user=config.fe_user,
                                     password=config.fe_password, http_port=config.fe_http_port)
     client.clean(database_name)
     client.create_database(database_name)
     client.use(database_name)
     return client


 def check2(client1, sql1, client2=None, sql2=None, forced=False):
     """
     check2
     client can be PaloClient or mysql.cursor
     """
     if client2 is None and sql2 is None:
         assert 0 == 1, 'the input of check2() is wrong'
     elif client2 is None:
         client2 = client1
     elif sql2 is None:
         sql2 = sql1
     ret1 = client1.execute(sql1)
     ret2 = client2.execute(sql2)
     util.check(ret1, ret2, forced)
     return True


 def get_explain_rollup(client, sql):
     """
     Get explain table
     """
     result = client.execute('EXPLAIN ' + sql)
     if result is None:
         return None
     rollup_flag = 'TABLE: '
     explain_rollup = list()
     for element in result:
         message = element[0].lstrip()
         profile = message.split(',')
         for msg in profile:
             if msg.startswith(rollup_flag):
                 LOG.info(L('explain', msg=msg))
                 table = msg[len(rollup_flag):].rstrip(' ')
                 pattern = re.compile(r'[(](.*?)[)]', re.S)
                 rollup = re.findall(pattern, table)
                 explain_rollup.append(rollup[0])
     return explain_rollup


 def assert_stop_routine_load(ret, client, stop_job=None, info=''):
     """assert, if not stop routine load """
     ret = client.show_routine_load(stop_job)
     if not ret and stop_job is not None:
         try:
             client.stop_routine_load(stop_job)
         except Exception as e:
             print(str(e))
     assert ret, info


 def wait_commit(client, routine_load_job_name, committed_expect_num, timeout=600):
     """wait task committed"""
     print('expect commited rows: %s\n' % committed_expect_num)
     while timeout > 0:
         ret = client.show_routine_load(routine_load_job_name,)
         routine_load_job = palo_job.RoutineLoadJob(ret[0])
         loaded_rows = routine_load_job.get_loaded_rows()
         print(loaded_rows)
         if str(loaded_rows) == str(committed_expect_num):
             return True
         state = routine_load_job.get_state()
         if state != 'RUNNING':
             print('routine load job\' state is not running, it\'s %s' % state)
             return False
         timeout -= 3
         time.sleep(3)
     return False


 def execute_ignore_error(func, *argc, **kwargs):
     """execute ignore sql error"""
     try:
         func(*argc, **kwargs)
     except Exception as e:
         LOG.info(L('execute ignore error', msg=str(e)))
         print(e)


 def execute_retry_when_msg(msg, func, *argc, **kwargs):
     """当SQL返回的结果包含msg，则重试"""
     retry = 20
     while retry > 0:
         try:
             ret = func(*argc, **kwargs)
             return ret
         except Exception as e:
             if msg in str(e):
                 time.sleep(3)
                 retry -= 1
             LOG.info(L("get error msg", msg=str(e)))
     return False


 def check_by_file(expect_file, table_name=None, sql=None,
                   client=None, database_name=None, **kwargs):
     """
     通过csv校验文件，验证表或sql的结果。
     根据sql执行返回的列的类型，读取csv文件转化为相应的结构，进行校验
     expect_file: 校验文件
     table_name：表名，被校验的表
     sql: sql，被校验的sql查询，表和sql不能同时为None
     client: 执行sql的client，如果为None则重新创建连接
     database_name: 数据库名称
     **kwargs：可以指定sql或表中的类型，可在复杂类型中使用，例： k2=palo_types.ARRAY_INT
               如array返回为字符串类型，对于array<double/fload>等可单独指定类型校验
     """
     LOG.info(L('check file', file=expect_file))
     if table_name is None and sql is None:
         assert 0 == 1, 'there is no table and query to be checked'
     if client is None:
         client = get_client()
         if database_name:
             client.use(database_name)
     if table_name:
         if database_name:
             table_name = '%s.%s' % (database_name, table_name)
         sql = 'select * from %s' % table_name
     cursor, ret = client.execute(sql, return_cursor=True)
     result_info = cursor.description
     column_name_list = util.get_attr(result_info, 0)
     column_type_list = util.get_attr(result_info, 1)
     for col_name, col_type in kwargs.items():
         try:
             idx = column_name_list.index(col_name)
             column_type_list[idx] = col_type
         except Exception as e:
             print(type(col_name), type(column_name_list[1]))
             print("%s is not the column in %s" % (col_name, column_name_list))
     expect_ret = palo_types.convert_csv_to_ret(expect_file, column_type_list)
     # 判断sql结果是否含有复杂类型，如果有，则需要对复杂类型进行处理
     if max(column_type_list) > 1000:
         ret = palo_types.convert_ret_complex_type(ret, column_type_list)
     util.check(ret, expect_ret, True)
     return True


 def check_by_sql(tested_sql, expect_sql, client=None, database_name=None, **kwargs):
     """
     适用于复杂类型的两个sql结果校验。将两个sql分别执行，格式化处理，保存到文件中进行校验
     通过**kwargs指定列的类型
     tested_sql: 被测sql
     expect_sql: 校验sql，预期正确的结果
     """
     if client is None:
         client = get_client()
         if database_name:
             client.use(database_name)

     cursor, expect_ret = client.execute(expect_sql, return_cursor=True)
     # 根据cursor获取返回结果的列名和类型code
     result_info = cursor.description
     column_name_list = util.get_attr(result_info, 0)
     column_type_list = util.get_attr(result_info, 1)
     tested_ret = client.execute(tested_sql)

     for col_name, col_type in kwargs.items():
         try:
             idx = column_name_list.index(col_name)
             column_type_list[idx] = col_type
         except Exception as e:
             print(type(col_name), type(column_name_list[1]))
             print("%s is not the column in %s" % (col_name, column_name_list))

     if len(kwargs) != 0:
         processed_expect_ret = palo_types.convert_ret_complex_type(expect_ret, column_type_list)
         processed_tested_ret = palo_types.convert_ret_complex_type(tested_ret, column_type_list)
         util.check(processed_tested_ret, processed_expect_ret, True)
     else:
         util.check(tested_ret, expect_ret, True)
     return True
	#!/bin/env python
	# -- coding: utf-8 --
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.

	###########################################################################
	# @file common.py
	# @date 2020-07-14 19:11:39
	# @brief common function
	############################################################################
	"""
	common test tool
	"""
	import re
	import time

	import palo_client
	import palo_config
	import palo_logger
	import palo_job
	import palo_types
	import util

	config = palo_config.config
	LOG = palo_logger.Logger.getLogger()
	L = palo_logger.StructedLogMessage


	def get_client(host=None):
	"""get a new client"""
	if host is None:
	host = config.fe_host
	client = palo_client.get_client(host, config.fe_query_port, user=config.fe_user,
	password=config.fe_password, http_port=config.fe_http_port)
	return client


	def create_workspace(database_name):
	"""get new palo client connect, and create db, and use and return client"""
	client = palo_client.get_client(config.fe_host, config.fe_query_port, user=config.fe_user,
	password=config.fe_password, http_port=config.fe_http_port)
	client.clean(database_name)
	client.create_database(database_name)
	client.use(database_name)
	return client


	def check2(client1, sql1, client2=None, sql2=None, forced=False):
	"""
	check2
	client can be PaloClient or mysql.cursor
	"""
	if client2 is None and sql2 is None:
	assert 0 == 1, 'the input of check2() is wrong'
	elif client2 is None:
	client2 = client1
	elif sql2 is None:
	sql2 = sql1
	ret1 = client1.execute(sql1)
	ret2 = client2.execute(sql2)
	util.check(ret1, ret2, forced)
	return True


	def get_explain_rollup(client, sql):
	"""
	Get explain table
	"""
	result = client.execute('EXPLAIN ' + sql)
	if result is None:
	return None
	rollup_flag = 'TABLE: '
	explain_rollup = list()
	for element in result:
	message = element[0].lstrip()
	profile = message.split(',')
	for msg in profile:
	if msg.startswith(rollup_flag):
	LOG.info(L('explain', msg=msg))
	table = msg[len(rollup_flag):].rstrip(' ')
	pattern = re.compile(r'[(](.*?)[)]', re.S)
	rollup = re.findall(pattern, table)
	explain_rollup.append(rollup[0])
	return explain_rollup


	def assert_stop_routine_load(ret, client, stop_job=None, info=''):
	"""assert, if not stop routine load """
	ret = client.show_routine_load(stop_job)
	if not ret and stop_job is not None:
	try:
	client.stop_routine_load(stop_job)
	except Exception as e:
	print(str(e))
	assert ret, info


	def wait_commit(client, routine_load_job_name, committed_expect_num, timeout=600):
	"""wait task committed"""
	print('expect commited rows: %s\n' % committed_expect_num)
	while timeout > 0:
	ret = client.show_routine_load(routine_load_job_name,)
	routine_load_job = palo_job.RoutineLoadJob(ret[0])
	loaded_rows = routine_load_job.get_loaded_rows()
	print(loaded_rows)
	if str(loaded_rows) == str(committed_expect_num):
	return True
	state = routine_load_job.get_state()
	if state != 'RUNNING':
	print('routine load job\' state is not running, it\'s %s' % state)
	return False
	timeout -= 3
	time.sleep(3)
	return False


	def execute_ignore_error(func, argc, *kwargs):
	"""execute ignore sql error"""
	try:
	func(argc, *kwargs)
	except Exception as e:
	LOG.info(L('execute ignore error', msg=str(e)))
	print(e)


	def execute_retry_when_msg(msg, func, argc, *kwargs):
	"""当SQL返回的结果包含msg，则重试"""
	retry = 20
	while retry > 0:
	try:
	ret = func(argc, *kwargs)
	return ret
	except Exception as e:
	if msg in str(e):
	time.sleep(3)
	retry -= 1
	LOG.info(L("get error msg", msg=str(e)))
	return False


	def check_by_file(expect_file, table_name=None, sql=None,
	client=None, database_name=None, **kwargs):
	"""
	通过csv校验文件，验证表或sql的结果。
	根据sql执行返回的列的类型，读取csv文件转化为相应的结构，进行校验
	expect_file: 校验文件
	table_name：表名，被校验的表
	sql: sql，被校验的sql查询，表和sql不能同时为None
	client: 执行sql的client，如果为None则重新创建连接
	database_name: 数据库名称
	**kwargs：可以指定sql或表中的类型，可在复杂类型中使用，例： k2=palo_types.ARRAY_INT
	如array返回为字符串类型，对于array<double/fload>等可单独指定类型校验
	"""
	LOG.info(L('check file', file=expect_file))
	if table_name is None and sql is None:
	assert 0 == 1, 'there is no table and query to be checked'
	if client is None:
	client = get_client()
	if database_name:
	client.use(database_name)
	if table_name:
	if database_name:
	table_name = '%s.%s' % (database_name, table_name)
	sql = 'select * from %s' % table_name
	cursor, ret = client.execute(sql, return_cursor=True)
	result_info = cursor.description
	column_name_list = util.get_attr(result_info, 0)
	column_type_list = util.get_attr(result_info, 1)
	for col_name, col_type in kwargs.items():
	try:
	idx = column_name_list.index(col_name)
	column_type_list[idx] = col_type
	except Exception as e:
	print(type(col_name), type(column_name_list[1]))
	print("%s is not the column in %s" % (col_name, column_name_list))
	expect_ret = palo_types.convert_csv_to_ret(expect_file, column_type_list)
	# 判断sql结果是否含有复杂类型，如果有，则需要对复杂类型进行处理
	if max(column_type_list) > 1000:
	ret = palo_types.convert_ret_complex_type(ret, column_type_list)
	util.check(ret, expect_ret, True)
	return True


	def check_by_sql(tested_sql, expect_sql, client=None, database_name=None, **kwargs):
	"""
	适用于复杂类型的两个sql结果校验。将两个sql分别执行，格式化处理，保存到文件中进行校验
	通过**kwargs指定列的类型
	tested_sql: 被测sql
	expect_sql: 校验sql，预期正确的结果
	"""
	if client is None:
	client = get_client()
	if database_name:
	client.use(database_name)

	cursor, expect_ret = client.execute(expect_sql, return_cursor=True)
	# 根据cursor获取返回结果的列名和类型code
	result_info = cursor.description
	column_name_list = util.get_attr(result_info, 0)
	column_type_list = util.get_attr(result_info, 1)
	tested_ret = client.execute(tested_sql)

	for col_name, col_type in kwargs.items():
	try:
	idx = column_name_list.index(col_name)
	column_type_list[idx] = col_type
	except Exception as e:
	print(type(col_name), type(column_name_list[1]))
	print("%s is not the column in %s" % (col_name, column_name_list))

	if len(kwargs) != 0:
	processed_expect_ret = palo_types.convert_ret_complex_type(expect_ret, column_type_list)
	processed_tested_ret = palo_types.convert_ret_complex_type(tested_ret, column_type_list)
	util.check(processed_tested_ret, processed_expect_ret, True)
	else:
	util.check(tested_ret, expect_ret, True)
	return True