def test_query_masking(self):
query_result = ResultSet()
new_engine = PgSQLEngine(instance=self.ins)
masking_result = new_engine.query_masking(db_name=0,
sql='',
resultset=query_result)
self.assertEqual(masking_result, query_result)
def test_query(self, _conn, _cursor, _execute):
_conn.return_value.cursor.return_value.fetchmany.return_value = [(1, )]
new_engine = PgSQLEngine(instance=self.ins)
query_result = new_engine.query(db_name=0,
sql='select 1',
limit_num=100)
self.assertIsInstance(query_result, ResultSet)
self.assertListEqual(query_result.rows, [(1, )])
def test_query_check_star_sql(self):
sql = "select * from xx "
new_engine = PgSQLEngine(instance=self.ins)
check_result = new_engine.query_check(db_name='archery', sql=sql)
self.assertDictEqual(
check_result, {
'msg': 'SQL语句中含有 * ',
'bad_query': False,
'filtered_sql': sql.strip(),
'has_star': True
})
def test_query_check_disable_sql(self):
sql = "update xxx set a=1 "
new_engine = PgSQLEngine(instance=self.ins)
check_result = new_engine.query_check(db_name='archery', sql=sql)
self.assertDictEqual(
check_result, {
'msg': '不支持的查询语法类型!',
'bad_query': True,
'filtered_sql': sql.strip(),
'has_star': False
})
def test_query_not_limit(self, _conn, _cursor, _execute):
_conn.return_value.cursor.return_value.fetchall.return_value = [(1, )]
new_engine = PgSQLEngine(instance=self.ins)
query_result = new_engine.query(db_name=0, sql='select 1', limit_num=0)
self.assertIsInstance(query_result, ResultSet)
def test_filter_sql_with_limit(self):
sql = "select * from xx limit 10"
new_engine = PgSQLEngine(instance=self.ins)
check_result = new_engine.filter_sql(sql=sql, limit_num=1)
self.assertEqual(check_result, "select * from xx limit 10;")
def test_describe_table(self, _query):
new_engine = PgSQLEngine(instance=self.ins)
describe = new_engine.describe_table(db_name='archery',
schema_name='archery',
tb_name='text')
self.assertIsInstance(describe, ResultSet)
def test_get_all_columns_by_tb(self, _query):
new_engine = PgSQLEngine(instance=self.ins)
columns = new_engine.get_all_columns_by_tb(db_name='archery',
tb_name='test2',
schema_name='archery')
self.assertListEqual(columns, ['id', 'name'])
def test_get_all_tables(self, _query):
new_engine = PgSQLEngine(instance=self.ins)
tables = new_engine.get_all_tables(db_name='archery',
schema_name='archery')
self.assertListEqual(tables, ['test2'])
def test_get_all_schemas(self, _query):
new_engine = PgSQLEngine(instance=self.ins)
schemas = new_engine.get_all_schemas(db_name='archery')
self.assertListEqual(schemas, ['archery'])
def test_get_all_databases(self, _query):
new_engine = PgSQLEngine(instance=self.ins)
dbs = new_engine.get_all_databases()
self.assertListEqual(dbs, ['archery'])
def test_get_connection(self, _conn):
new_engine = PgSQLEngine(instance=self.ins)
new_engine.get_connection()
_conn.assert_called_once()
