From 69e29b11d8ea95184558015996b54fa5e634c4cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E6=99=93=E9=A3=9E?= <feiazifeiazi@hotmail.com>
Date: Fri, 27 Dec 2024 12:38:05 +0800
Subject: [PATCH] =?UTF-8?q?Mongo=E8=BF=94=E5=9B=9E=E7=BB=93=E6=9E=9C?=
 =?UTF-8?q?=E6=94=AF=E6=8C=81=E8=84=B1=E6=95=8F-V0.1-beta=20(#2884)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 返回一个脱敏后的结果集

* 通过goInception获取select list

* test_query_masking

* test_query_masking

* 8889

* 18888888889
---
 sql/engines/mongo.py      |  8 ++++++++
 sql/engines/tests.py      | 33 ++++++++++++++++++++++++++++++++-
 sql/utils/data_masking.py | 23 ++++++++++++++++++-----
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/sql/engines/mongo.py b/sql/engines/mongo.py
index 4214a03845..b7c879d4bf 100644
--- a/sql/engines/mongo.py
+++ b/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
diff --git a/sql/engines/tests.py b/sql/engines/tests.py
index 8cb5c63eb3..cd2ae7bba4 100644
--- a/sql/engines/tests.py
+++ b/sql/engines/tests.py
@@ -17,7 +17,7 @@
 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
+from sql.models import DataMaskingColumns, Instance, SqlWorkflow, SqlWorkflowContent
 
 User = get_user_model()
 
@@ -1665,9 +1665,19 @@ def setUp(self) -> None:
         )
         self.engine = MongoEngine(instance=self.ins)
         self.sys_config = SysConfig()
+        # rule_type=100的规则不需要加，会自动创建。只需要加脱敏字段
+        DataMaskingColumns.objects.create(
+            rule_type=100,
+            active=True,
+            instance=self.ins,
+            table_schema="*",
+            table_name="*",
+            column_name="mobile",
+        )
 
     def tearDown(self) -> None:
         self.ins.delete()
+        DataMaskingColumns.objects.all().delete()
 
     @patch("sql.engines.mongo.pymongo")
     def test_get_connection(self, mock_pymongo):
@@ -2042,6 +2052,27 @@ def test_create_instance_user(self, _mock_command):
             ],
         )
 
+    def test_query_masking(self):
+        query_result = ResultSet()
+        new_engine = MongoEngine(instance=self.ins)
+        query_result.column_list = ["id", "mobile"]
+        query_result.rows = (
+            ("a11", "18888888888"),
+            ("a12", ""),
+            ("a13", None),
+            ("a14", "18888888889"),
+        )
+        masking_result = new_engine.query_masking(
+            db_name="archery", sql="db.test_collection.find()", resultset=query_result
+        )
+        mask_result_rows = [
+            ["a11", "188****8888"],
+            ["a12", ""],
+            ["a13", None],
+            ["a14", "188****8889"],
+        ]
+        self.assertEqual(masking_result.rows, mask_result_rows)
+
 
 class TestClickHouse(TestCase):
     def setUp(self):
diff --git a/sql/utils/data_masking.py b/sql/utils/data_masking.py
index 61ad93c9c1..08e7c63979 100644
--- a/sql/utils/data_masking.py
+++ b/sql/utils/data_masking.py
@@ -24,11 +24,24 @@ 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:
+            # 通过goInception获取select list
+            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