返回一个脱敏后的结果集

sql/engines/mongo.py

@@ -15,6 +15,8 @@
 from bson.objectid import ObjectId
 from bson.int64 import Int64
 
+from sql.utils.data_masking import data_masking
+
 from . import EngineBase
 from .models import ResultSet, ReviewSet, ReviewResult
 from common.config import SysConfig
@@ -1422,3 +1424,9 @@ def reset_instance_user_pwd(self, db_name_user: str, reset_pwd: str, **kwargs):
         except Exception as e:
             exec_result.error = str(e)
         return exec_result
+
+    def query_masking(self, db_name=None, sql="", resultset=None):
+        """传入 sql语句, db名, 结果集,
+        返回一个脱敏后的结果集"""
+        mask_result = data_masking(self.instance, db_name, sql, resultset)
+        return mask_result

sql/utils/data_masking.py

@@ -24,11 +24,23 @@ def data_masking(instance, db_name, sql, sql_result):
         for token in p.tokens:
             if token.ttype is Keyword and token.value.upper() in ["UNION", "UNION ALL"]:
                 keywords_count["UNION"] = keywords_count.get("UNION", 0) + 1
-        # 通过goInception获取select list
-        inception_engine = GoInceptionEngine()
-        select_list = inception_engine.query_data_masking(
-            instance=instance, db_name=db_name, sql=sql
-        )
+        if instance.db_type == "mongo":
+            select_list = [
+                {
+                    "index": index,
+                    "field": field,
+                    "type": "varchar",
+                    "table": "*",
+                    "schema": db_name,
+                    "alias": field,
+                }
+                for index, field in enumerate(sql_result.column_list)
+            ]
+        else:
+            inception_engine = GoInceptionEngine()
+            select_list = inception_engine.query_data_masking(
+                instance=instance, db_name=db_name, sql=sql
+            )
         # 如果UNION存在，那么调用去重函数
         select_list = (
             del_repeat(select_list, keywords_count) if keywords_count else select_list