def _query_apply_audit_call_back(apply_id, workflow_status):
"""
查询权限申请用于工作流审核回调
:param apply_id: 申请id
:param workflow_status: 审核结果
:return:
"""
# 更新业务表状态
apply_info = QueryPrivilegesApply.objects.get(apply_id=apply_id)
apply_info.status = workflow_status
apply_info.save()
# 审核通过插入权限信息，批量插入，减少性能消耗
if workflow_status == WorkflowDict.workflow_status['audit_success']:
apply_queryset = QueryPrivilegesApply.objects.get(apply_id=apply_id)
# 库权限
if apply_queryset.priv_type == 1:
insert_list = [
QueryPrivileges(user_name=apply_queryset.user_name,
user_display=apply_queryset.user_display,
instance=apply_queryset.instance,
db_name=db_name,
table_name=apply_queryset.table_list,
valid_date=apply_queryset.valid_date,
limit_num=apply_queryset.limit_num,
priv_type=apply_queryset.priv_type)
for db_name in apply_queryset.db_list.split(',')
]
# 表权限
elif apply_queryset.priv_type == 2:
insert_list = [
QueryPrivileges(user_name=apply_queryset.user_name,
user_display=apply_queryset.user_display,
instance=apply_queryset.instance,
db_name=apply_queryset.db_list,
table_name=table_name,
valid_date=apply_queryset.valid_date,
limit_num=apply_queryset.limit_num,
priv_type=apply_queryset.priv_type)
for table_name in apply_queryset.table_list.split(',')
]
QueryPrivileges.objects.bulk_create(insert_list)
def setUp(self):
self.slave1 = Instance(instance_name='test_slave_instance', type='slave', db_type='mysql',
host='testhost', port=3306, user='******', password='******')
self.slave2 = Instance(instance_name='test_instance_non_mysql', type='slave', db_type='mssql',
host='some_host2', port=3306, user='******', password='******')
self.slave1.save()
self.slave2.save()
archer_user = get_user_model()
self.superuser1 = archer_user(username='******', is_superuser=True)
self.superuser1.save()
self.u1 = archer_user(username='******', display='中文显示', is_active=True)
self.u1.save()
self.u2 = archer_user(username='******', display='中文显示', is_active=True)
self.u2.save()
self.u3 = archer_user(username='******', display='中文显示', is_active=True)
self.u3.save()
sql_query_perm = Permission.objects.get(codename='query_submit')
self.u2.user_permissions.add(sql_query_perm)
self.u3.user_permissions.add(sql_query_perm)
tomorrow = datetime.now() + timedelta(days=1)
self.query_apply_1 = QueryPrivilegesApply(
group_id=1,
group_name='some_group',
title='some_title',
user_name='some_user',
instance_name='some_ins',
db_list='some_db,some_db2',
limit_num=100,
valid_date=tomorrow,
priv_type=1,
status=0,
audit_auth_groups='some_audit_group'
)
self.query_apply_1.save()
self.query_apply_2 = QueryPrivilegesApply(
group_id=1,
group_name='some_group',
title='some_title',
user_name='some_user',
instance_name='some_ins',
db_list='some_db',
table_list='some_table,some_tb2',
limit_num=100,
valid_date=tomorrow,
priv_type=2,
status=0,
audit_auth_groups='some_audit_group'
)
self.query_apply_2.save()
self.db_priv_for_user3 = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave1.instance_name,
db_name='some_db',
table_name='',
valid_date=tomorrow,
limit_num=70,
priv_type=1)
self.db_priv_for_user3.save()
self.table_priv_for_user3 = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave1.instance_name,
db_name='another_db',
table_name='some_table',
valid_date=tomorrow,
limit_num=60,
priv_type=2)
self.table_priv_for_user3.save()
self.db_priv_for_user3_another_instance = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave2.instance_name,
db_name='some_db_another_instance',
table_name='',
valid_date=tomorrow,
limit_num=50,
priv_type=1)
self.db_priv_for_user3_another_instance.save()
class QueryTest(TestCase):
def setUp(self):
self.slave1 = Instance(instance_name='test_slave_instance', type='slave', db_type='mysql',
host='testhost', port=3306, user='******', password='******')
self.slave2 = Instance(instance_name='test_instance_non_mysql', type='slave', db_type='mssql',
host='some_host2', port=3306, user='******', password='******')
self.slave1.save()
self.slave2.save()
archer_user = get_user_model()
self.superuser1 = archer_user(username='******', is_superuser=True)
self.superuser1.save()
self.u1 = archer_user(username='******', display='中文显示', is_active=True)
self.u1.save()
self.u2 = archer_user(username='******', display='中文显示', is_active=True)
self.u2.save()
self.u3 = archer_user(username='******', display='中文显示', is_active=True)
self.u3.save()
sql_query_perm = Permission.objects.get(codename='query_submit')
self.u2.user_permissions.add(sql_query_perm)
self.u3.user_permissions.add(sql_query_perm)
tomorrow = datetime.now() + timedelta(days=1)
self.query_apply_1 = QueryPrivilegesApply(
group_id=1,
group_name='some_group',
title='some_title',
user_name='some_user',
instance_name='some_ins',
db_list='some_db,some_db2',
limit_num=100,
valid_date=tomorrow,
priv_type=1,
status=0,
audit_auth_groups='some_audit_group'
)
self.query_apply_1.save()
self.query_apply_2 = QueryPrivilegesApply(
group_id=1,
group_name='some_group',
title='some_title',
user_name='some_user',
instance_name='some_ins',
db_list='some_db',
table_list='some_table,some_tb2',
limit_num=100,
valid_date=tomorrow,
priv_type=2,
status=0,
audit_auth_groups='some_audit_group'
)
self.query_apply_2.save()
self.db_priv_for_user3 = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave1.instance_name,
db_name='some_db',
table_name='',
valid_date=tomorrow,
limit_num=70,
priv_type=1)
self.db_priv_for_user3.save()
self.table_priv_for_user3 = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave1.instance_name,
db_name='another_db',
table_name='some_table',
valid_date=tomorrow,
limit_num=60,
priv_type=2)
self.table_priv_for_user3.save()
self.db_priv_for_user3_another_instance = QueryPrivileges(
user_name=self.u3.username,
user_display=self.u3.display,
instance_name=self.slave2.instance_name,
db_name='some_db_another_instance',
table_name='',
valid_date=tomorrow,
limit_num=50,
priv_type=1)
self.db_priv_for_user3_another_instance.save()
def tearDown(self):
self.u1.delete()
self.u2.delete()
self.u3.delete()
self.slave1.delete()
self.slave2.delete()
self.query_apply_1.delete()
self.query_apply_2.delete()
QueryPrivileges.objects.all().delete()
archer_config = SysConfig()
archer_config.set('disable_star', False)
def testQueryAuditCallback(self):
"""测试权限申请工单回调"""
# 工单状态改为审核失败, 验证工单状态
query.query_audit_call_back(self.query_apply_1.apply_id, 2)
self.query_apply_1.refresh_from_db()
self.assertEqual(self.query_apply_1.status, 2)
for db in self.query_apply_1.db_list.split(','):
self.assertEqual(len(QueryPrivileges.objects.filter(
user_name=self.query_apply_1.user_name,
db_name=db,
limit_num=100)), 0)
# 工单改为审核成功, 验证工单状态和权限状态
query.query_audit_call_back(self.query_apply_1.apply_id, 1)
self.query_apply_1.refresh_from_db()
self.assertEqual(self.query_apply_1.status, 1)
for db in self.query_apply_1.db_list.split(','):
self.assertEqual(len(QueryPrivileges.objects.filter(
user_name=self.query_apply_1.user_name,
db_name=db,
limit_num=100)), 1)
# 表权限申请测试, 只测试审核成功
query.query_audit_call_back(self.query_apply_2.apply_id, 1)
self.query_apply_2.refresh_from_db()
self.assertEqual(self.query_apply_2.status, 1)
for tb in self.query_apply_2.table_list.split(','):
self.assertEqual(len(QueryPrivileges.objects.filter(
user_name=self.query_apply_2.user_name,
db_name=self.query_apply_2.db_list,
table_name=tb,
limit_num=self.query_apply_2.limit_num)), 1)
@patch('sql.engines.mysql.MysqlEngine.query')
@patch('sql.engines.mysql.MysqlEngine.query_masking')
@patch('sql.query.query_priv_check')
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'])
@patch('sql.engines.mysql.MysqlEngine.query')
@patch('sql.engines.mysql.MysqlEngine.query_masking')
@patch('sql.query.query_priv_check')
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)
@patch('sql.query.query_priv_check')
def testStarOptionOn(self, _priv_check):
c = Client()
c.force_login(self.u2)
some_limit = 100
sql_with_star = 'select * from some_table'
some_db = 'some_db'
_priv_check.return_value = {'status': 0, 'data': {'limit_num': 100, 'priv_check': 1}}
archer_config = SysConfig()
archer_config.set('disable_star', True)
r = c.post('/query/', data={'instance_name': self.slave1.instance_name,
'sql_content': sql_with_star,
'db_name': some_db,
'limit_num': some_limit})
archer_config.set('disable_star', False)
r_json = r.json()
self.assertEqual(1, r_json['status'])
@patch('sql.query.Masking')
def test_query_priv_check(self, mock_masking):
# 超级用户直接返回
superuser_limit = query.query_priv_check(self.superuser1, self.slave1.instance_name,
'some_db', 'some_sql', 100)
self.assertEqual(superuser_limit['status'], 0)
self.assertEqual(superuser_limit['data']['limit_num'], 100)
# 无语法树解析，只校验db_name
limit_without_tree_analyse = query.query_priv_check(self.u3, self.slave2,
'some_db_another_instance','some_sql', 1000)
self.assertEqual(limit_without_tree_analyse['data']['limit_num'], self.db_priv_for_user3_another_instance.limit_num)
# 无语法树解析， 无权限的情况
limit_without_tree_analyse = query.query_priv_check(self.u3, self.slave2,
'some_db_does_not_exist','some_sql', 1000)
self.assertEqual(limit_without_tree_analyse['status'],1)
self.assertIn('some_db_does_not_exist', limit_without_tree_analyse['msg'])
# 有语法树解析, 有库权限
mock_masking.return_value.query_table_ref.return_value = {
'status': 0,
'data': [{
'db': 'another_db',
'table': 'some_table'
}]}
limit_with_tree_analyse = query.query_priv_check(self.u3, self.slave1,
'some_db','some_sql',1000)
mock_masking.return_value.query_table_ref.assert_called_once()
self.assertEqual(limit_with_tree_analyse['data']['limit_num'], self.table_priv_for_user3.limit_num)