def testSQLWithoutLimit(self, _priv_check, _query_masking, _query):
c = Client()
some_limit = 100
sql_without_limit = 'select some from some_table'
sql_with_limit = 'select some from some_table limit {0};'.format(some_limit)
some_db = 'some_db'
c.force_login(self.u2)
q_result = ResultSet(full_sql=sql_without_limit, rows=['value'])
q_result.column_list = ['some']
_query.return_value = q_result
_query_masking.return_value = q_result
_priv_check.return_value = {'status': 0, 'data': {'limit_num': 100, 'priv_check': 1}}
r = c.post('/query/', data={'instance_name': self.slave1.instance_name,
'sql_content': sql_without_limit,
'db_name': some_db,
'limit_num': some_limit})
_query.assert_called_once_with(db_name=some_db, sql=sql_with_limit, limit_num=some_limit)
r_json = r.json()
self.assertEqual(r_json['data']['rows'], ['value'])
self.assertEqual(r_json['data']['column_list'], ['some'])
# 带 * 且不带 limit 的sql
sql_with_star = 'select * from some_table'
filtered_sql_with_star = 'select * from some_table limit {0};'.format(some_limit)
_query.reset_mock()
c.post('/query/', data={'instance_name': self.slave1.instance_name,
'sql_content': sql_with_star,
'db_name': some_db,
'limit_num': some_limit})
_query.assert_called_once_with(db_name=some_db, sql=filtered_sql_with_star, limit_num=some_limit)
def testCorrectSQL(self, _priv_check, _query_masking, _query):
c = Client()
some_sql = 'select some from some_table limit 100;'
some_db = 'some_db'
some_limit = 100
c.force_login(self.u1)
r = c.post('/query/', data={'instance_name': self.slave1.instance_name,
'sql_content': some_sql,
'db_name': some_db,
'limit_num': some_limit})
self.assertEqual(r.status_code, 403)
c.force_login(self.u2)
q_result = ResultSet(full_sql=some_sql, rows=['value'])
q_result.column_list = ['some']
_query.return_value = q_result
_query_masking.return_value = q_result
_priv_check.return_value = {'status': 0, 'data': {'limit_num': 100, 'priv_check': 1}}
r = c.post('/query/', data={'instance_name': self.slave1.instance_name,
'sql_content': some_sql,
'db_name': some_db,
'limit_num': some_limit})
_query.assert_called_once_with(db_name=some_db, sql=some_sql, limit_num=some_limit)
r_json = r.json()
self.assertEqual(r_json['data']['rows'], ['value'])
self.assertEqual(r_json['data']['column_list'], ['some'])
