修复distinct关键字的报错

pengkang1991 · Nov 29, 2023 · 0898004 · 0898004
1 parent e40d728
commit 0898004
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/sql/engines/mssql.py b/sql/engines/mssql.py
@@ -316,6 +316,10 @@ def filter_sql(self, sql="", limit_num=0):
         # 对查询sql增加limit限制
         if re.match(r"^select", sql_lower):
             if sql_lower.find(" top ") == -1:
+                if sql_lower.find(" distinct ") > 0:
+                    return sql_lower.replace(
+                        "distinct", "distinct top {}".format(limit_num)
+                    )
                 return sql_lower.replace("select", "select top {}".format(limit_num))
         return sql.strip()
 

diff --git a/sql/engines/tests.py b/sql/engines/tests.py
@@ -192,6 +192,13 @@ def test_filter_sql(self):
         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 = (