def testAllColumns(self, mock_query):
db_result = ResultSet()
db_result.rows = [('col_1', 'type'), ('col_2', 'type2')]
mock_query.return_value = db_result
new_engine = MysqlEngine(instance=self.ins1)
dbs = new_engine.get_all_columns_by_tb('some_db', 'some_tb')
self.assertEqual(dbs, ['col_1', 'col_2'])
def test_query_check_update_sql(self):
new_engine = MysqlEngine(instance=self.ins1)
update_sql = 'update user set id=0'
check_result = new_engine.query_check(db_name='some_db', sql=update_sql)
self.assertDictEqual(check_result,
{'msg': '不支持的查询语法类型!', 'bad_query': True, 'filtered_sql': 'update user set id=0',
'has_star': False})
def test_execute_workflow(self, _inception_engine):
sql = 'update user set id=1'
_inception_engine.return_value.execute.return_value = ReviewSet(
full_sql=sql)
new_engine = MysqlEngine(instance=self.ins1)
execute_result = new_engine.execute_workflow(self.wf)
self.assertIsInstance(execute_result, ReviewSet)
def test_query_masking_not_select(self, _data_masking):
query_result = ResultSet()
new_engine = MysqlEngine(instance=self.ins1)
masking_result = new_engine.query_masking(db_name='archery',
sql='explain select 1',
resultset=query_result)
self.assertEqual(masking_result, query_result)
def test_execute_check_normal_sql_with_Exception(self, _inception_engine):
sql = 'update user set id=1'
_inception_engine.return_value.execute_check.side_effect = RuntimeError(
)
new_engine = MysqlEngine(instance=self.ins1)
with self.assertRaises(RuntimeError):
new_engine.execute_check(db_name=0, sql=sql)
def test_query_masking(self, _data_masking):
query_result = ResultSet()
new_engine = MysqlEngine(instance=self.ins1)
masking_result = new_engine.query_masking(db_name='archery',
sql='select 1',
resultset=query_result)
self.assertIsInstance(masking_result, ResultSet)
def testQueryCheck(self):
new_engine = MysqlEngine(instance=self.ins1)
sql_without_limit = '-- 测试\n select user from usertable'
check_result = new_engine.query_check(db_name='some_db',
sql=sql_without_limit)
self.assertEqual(check_result['filtered_sql'],
'select user from usertable')
def testAllDb(self, mock_query):
db_result = ResultSet()
db_result.rows = [('db_1',), ('db_2',)]
mock_query.return_value = db_result
new_engine = MysqlEngine(instance=self.ins1)
dbs = new_engine.get_all_databases()
self.assertEqual(dbs, ['db_1', 'db_2'])
def testAllTables(self, mock_query):
table_result = ResultSet()
table_result.rows = [('tb_1', 'some_des'), ('tb_2', 'some_des')]
mock_query.return_value = table_result
new_engine = MysqlEngine(instance=self.ins1)
tables = new_engine.get_all_tables('some_db')
mock_query.assert_called_once_with(db_name='some_db', sql=ANY)
self.assertEqual(tables, ['tb_1', 'tb_2'])
def test_execute_check_select_sql(self):
sql = 'select * from user'
row = ReviewResult(id=1, errlevel=2,
stagestatus='驳回高危SQL',
errormessage='仅支持DML和DDL语句，查询语句请使用SQL查询功能！',
sql=sql)
new_engine = MysqlEngine(instance=self.ins1)
check_result = new_engine.execute_check(db_name='archery', sql=sql)
self.assertIsInstance(check_result, ReviewSet)
self.assertEqual(check_result.rows[0].__dict__, row.__dict__)
def test_query_check_wrong_sql(self):
new_engine = MysqlEngine(instance=self.ins1)
wrong_sql = '-- 测试'
check_result = new_engine.query_check(db_name='some_db', sql=wrong_sql)
self.assertDictEqual(
check_result, {
'msg': '不支持的查询语法类型!',
'bad_query': True,
'filtered_sql': '-- 测试',
'has_star': True
})
def test_execute_check_critical_sql(self):
self.sys_config.set('critical_ddl_regex', '^|update')
self.sys_config.get_all_config()
sql = 'update user set id=1'
row = ReviewResult(id=1, errlevel=2,
stagestatus='驳回高危SQL',
errormessage='禁止提交匹配' + '^|update' + '条件的语句！',
sql=sql)
new_engine = MysqlEngine(instance=self.ins1)
check_result = new_engine.execute_check(db_name='archery', sql=sql)
self.assertIsInstance(check_result, ReviewSet)
self.assertEqual(check_result.rows[0].__dict__, row.__dict__)
def testQuery(self, connect):
cur = Mock()
connect.return_value.cursor = cur
cur.return_value.execute = Mock()
cur.return_value.fetchmany.return_value = (('v1', 'v2'),)
cur.return_value.description = (('k1', 'some_other_des'), ('k2', 'some_other_des'))
new_engine = MysqlEngine(instance=self.ins1)
query_result = new_engine.query(sql='some_str', limit_num=100)
cur.return_value.execute.assert_called()
cur.return_value.fetchmany.assert_called_once_with(size=100)
connect.return_value.close.assert_called_once()
self.assertIsInstance(query_result, ResultSet)
def test_execute_check_normal_sql(self, _inception_engine):
sql = 'update user set id=1'
row = ReviewResult(id=1,
errlevel=0,
stagestatus='Audit completed',
errormessage='None',
sql=sql,
affected_rows=0,
execute_time=0, )
_inception_engine.return_value.execute_check.return_value = ReviewSet(full_sql=sql, rows=[row])
new_engine = MysqlEngine(instance=self.ins1)
check_result = new_engine.execute_check(db_name='archery', sql=sql)
self.assertIsInstance(check_result, ReviewSet)
self.assertEqual(check_result.rows[0].__dict__, row.__dict__)
def test_execute(self, _connect, _cursor, _execute):
new_engine = MysqlEngine(instance=self.ins1)
execute_result = new_engine.execute(self.wf)
self.assertIsInstance(execute_result, ResultSet)
def testDescribe(self, mock_query):
new_engine = MysqlEngine(instance=self.ins1)
new_engine.describe_table('some_db', 'some_db')
mock_query.assert_called_once()
def testGetConnection(self, connect):
new_engine = MysqlEngine(instance=self.ins1)
new_engine.get_connection()
connect.assert_called_once()
def test_set_variable(self, _query):
new_engine = MysqlEngine(instance=self.ins1)
new_engine.set_variable('binlog_format', 'ROW')
_query.assert_called_once()
def test_get_variables_filter(self, _query):
new_engine = MysqlEngine(instance=self.ins1)
new_engine.get_variables(variables=['binlog_format'])
_query.assert_called()
def test_get_variables_not_filter(self, _query):
new_engine = MysqlEngine(instance=self.ins1)
new_engine.get_variables()
_query.assert_called_once()
def test_filter_sql(self):
new_engine = MysqlEngine(instance=self.ins1)
sql_without_limit = 'select user from usertable'
check_result = new_engine.filter_sql(sql=sql_without_limit, limit_num=100)
self.assertEqual(check_result, 'select user from usertable limit 100;')
def test_server_version(self, _query):
_query.return_value.rows = (('5.7.20', ), )
new_engine = MysqlEngine(instance=self.ins1)
server_version = new_engine.server_version
self.assertTupleEqual(server_version, (5, 7, 20))