sql-使用python开发的sql查询平台-优质项目.zip

import json from datetime import timedelta, datetime from unittest.mock import patch, Mock, ANY import sqlparse from django.contrib.auth import get_user_model from django.test import TestCase from common.config import SysConfig from sql.engines import EngineBase from sql.engines.goinception import GoInceptionEngine from sql.engines.models import ResultSet, ReviewSet, ReviewResult from sql.engines.mssql import MssqlEngine from sql.engines.redis import RedisEngine from sql.engines.pgsql import PgSQLEngine from sql.engines.oracle import OracleEngine from sql.engines.mongo import MongoEngine from sql.engines.clickhouse import ClickHouseEngine from sql.engines.odps import ODPSEngine from sql.models import Instance, SqlWorkflow, SqlWorkflowContent User = get_user_model() class TestReviewSet(TestCase): def test_review_set(self): new_review_set = ReviewSet() new_review_set.rows = [{"id": "1679123"}] self.assertIn("1679123", new_review_set.json()) class TestEngineBase(TestCase): @classmethod def setUpClass(cls): cls.u1 = User(username="some_user", display="用户1") cls.u1.save() cls.ins1 = Instance( instance_name="some_ins", type="master", db_type="mssql", host="some_host", port=1366, user="ins_user", password="some_str", ) cls.ins1.save() cls.wf1 = SqlWorkflow.objects.create( workflow_name="some_name", group_id=1, group_name="g1", engineer=cls.u1.username, engineer_display=cls.u1.display, audit_auth_groups="some_group", create_time=datetime.now() - timedelta(days=1), status="workflow_finish", is_backup=True, instance=cls.ins1, db_name="some_db", syntax_type=1, ) cls.wfc1 = SqlWorkflowContent.objects.create( workflow=cls.wf1, sql_content="some_sql", execute_result=json.dumps([{"id": 1, "sql": "some_content"}]), ) @classmethod def tearDownClass(cls): cls.wfc1.delete() cls.wf1.delete() cls.ins1.delete() cls.u1.delete() def test_init_with_ins(self): engine = EngineBase(instance=self.ins1) self.assertEqual(self.ins1.instance_name, engine.instance_name) self.assertEqual(self.ins1.user, engine.user) class TestMssql(TestCase): @classmethod def setUpClass(cls): cls.ins1 = Instance( instance_name="some_ins", type="slave", db_type="mssql", host="some_host", port=1366, user="ins_user", password="some_str", ) cls.ins1.save() cls.engine = MssqlEngine(instance=cls.ins1) cls.wf = SqlWorkflow.objects.create( workflow_name="some_name", group_id=1, group_name="g1", engineer_display="", audit_auth_groups="some_group", create_time=datetime.now() - timedelta(days=1), status="workflow_finish", is_backup=True, instance=cls.ins1, db_name="some_db", syntax_type=1, ) SqlWorkflowContent.objects.create( workflow=cls.wf, sql_content="insert into some_tb values (1)" ) @classmethod def tearDownClass(cls): cls.ins1.delete() cls.wf.delete() SqlWorkflowContent.objects.all().delete() @patch("sql.engines.mssql.pyodbc.connect") def testGetConnection(self, connect): new_engine = MssqlEngine(instance=self.ins1) new_engine.get_connection() connect.assert_called_once() @patch("sql.engines.mssql.pyodbc.connect") 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 = MssqlEngine(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(100) connect.return_value.close.assert_called_once() self.assertIsInstance(query_result, ResultSet) @patch.object(MssqlEngine, "query") def testAllDb(self, mock_query): db_result = ResultSet() db_result.rows = [("db_1",), ("db_2",)] mock_query.return_value = db_result new_engine = MssqlEngine(instance=self.ins1) dbs = new_engine.get_all_databases() self.assertEqual(dbs.rows, ["db_1", "db_2"]) @patch.object(MssqlEngine, "query") 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 = MssqlEngine(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.rows, ["tb_1", "tb_2"]) @patch.object(MssqlEngine, "query") 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 = MssqlEngine(instance=self.ins1) dbs = new_engine.get_all_columns_by_tb("some_db", "some_tb") self.assertEqual(dbs.rows, ["col_1", "col_2"]) @patch.object(MssqlEngine, "query") def testDescribe(self, mock_query): new_engine = MssqlEngine(instance=self.ins1) new_engine.describe_table("some_db", "some_db") mock_query.assert_called_once() def testQueryCheck(self): new_engine = MssqlEngine(instance=self.ins1) # 只抽查一个函数 banned_sql = "select concat(phone,1) from user_table" check_result = new_engine.query_check(db_name="some_db", sql=banned_sql) self.assertTrue(check_result.get("bad_query")) banned_sql = "select phone from user_table where phone=concat(phone,1)" check_result = new_engine.query_check(db_name="some_db", sql=banned_sql) self.assertTrue(check_result.get("bad_query")) sp_sql = "sp_helptext '[SomeName].[SomeAction]'" check_result = new_engine.query_check(db_name="some_db", sql=sp_sql) self.assertFalse(check_result.get("bad_query")) self.assertEqual(check_result.get("filtered_sql"), sp_sql) def test_filter_sql(self): new_engine = MssqlEngine(instance=self.ins1) # 只抽查一个函数 banned_sql = "select user from user_table" check_result = new_engine.filter_sql(sql=banned_sql, limit_num=10) self.assertEqual(check_result, "select top 10 user from user_table") def test_filter_sql_with_distinct(self): new_engine = MssqlEngine(instance=self.ins1) # 只抽查一个函数 banned_sql = "select distinct * from user_table" check_result = new_engine.filter_sql(sql=banned_sql, limit_num=10) self.assertEqual(check_result, "select distinct top 10 * from user_table") def test_execute_check(self): new_engine = MssqlEngine(instance=self.ins1) test_sql = ( "use database\ngo\nsome sql1\nGO\nsome sql2\n\r\nGo\nsome sql3\n\r\ngO\n" ) check_result = new_engine.execute_check(db_name=None, sql=test_sql) self.assertIsInstance(check_result, ReviewSet) self.assertEqual(check_result.rows[1].__dict__["sql"], "use database\n") self.assertEqual(check_result.rows[2].__dict__["sql"], "\nsome sql1\n") self.assertEqual(check_result.rows[4].__dict__["sql"], "\nsome sql3\