添加用例数据SQL查询类型支持 (#174)

wu-clan · Apr 8, 2024 · 1b621a3 · 1b621a3
1 parent 09744ba
commit 1b621a3
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 9 deletions.
diff --git a/httpfpt/common/send_request.py b/httpfpt/common/send_request.py
@@ -169,8 +169,10 @@ def send_request(
                                 if relate_parsed_data:
                                     parsed_data = relate_parsed_data
                             elif key == SetupType.SQL:
-                                sql = var_extractor.vars_replace({'sql': value}, parsed_data['env'])['sql']
-                                mysql_client.exec_case_sql(sql, parsed_data['env'])
+                                setup_sql = var_extractor.vars_replace({'sql': value}, parsed_data['env'])
+                                sql = setup_sql['sql']
+                                sql_fetch = setup_sql.get('fetch')
+                                mysql_client.exec_case_sql(sql, sql_fetch, parsed_data['env'])
                             elif key == SetupType.HOOK:
                                 hook_executor.exec_hook_func(value)
                             elif key == SetupType.WAIT_TIME:
@@ -261,8 +263,10 @@ def send_request(
                     for key, value in item.items():
                         if value is not None:
                             if key == TeardownType.SQL:
-                                sql = var_extractor.vars_replace({'sql': value}, parsed_data['env'])['sql']
-                                mysql_client.exec_case_sql(sql, parsed_data['env'])
+                                teardown_sql = var_extractor.vars_replace({'sql': value}, parsed_data['env'])
+                                sql = teardown_sql['sql']
+                                sql_fetch = teardown_sql.get('fetch')
+                                mysql_client.exec_case_sql(sql, sql_fetch, parsed_data['env'])
                             if key == TeardownType.HOOK:
                                 hook_executor.exec_hook_func(value)
                             if key == TeardownType.EXTRACT:

diff --git a/httpfpt/db/mysql_db.py b/httpfpt/db/mysql_db.py
@@ -128,19 +128,22 @@ def execute(self, sql: str) -> int:
         finally:
             self.close(conn, cursor)
 
-    def exec_case_sql(self, sql: str, env: str | None = None) -> dict | list | int | None:
+    def exec_case_sql(self, sql: str, fetch: QueryFetchType | None, env: str | None = None) -> dict | list | int | None:
         """
         执行用例 sql
 
         :param sql:
+        :param fetch:
         :param env:
         :return:
         """
         # 获取返回数据
         if isinstance(sql, str):
             log.info(f'执行 SQL: {sql}')
             if sql.startswith(SqlType.select):
-                return self.query(sql)
+                if fetch is None:
+                    fetch = QueryFetchType.ALL
+                return self.query(sql, fetch)
             else:
                 return self.execute(sql)
 

diff --git a/httpfpt/schemas/case_data.py b/httpfpt/schemas/case_data.py
@@ -2,11 +2,14 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field
 from typing_extensions import Literal
 
+if TYPE_CHECKING:
+    from httpfpt.enums.query_fetch_type import QueryFetchType
+
 __all__ = [
     'CaseData',
     'CaseCacheData',
@@ -65,8 +68,9 @@ class SetupTestCaseData(BaseModel):
 
 class SetupSqlData(BaseModel):
     key: str
-    type: Literal['cache', 'wnv', 'global']
+    type: Literal['cache', 'env', 'global']
     sql: str
+    fetch: QueryFetchType | None = None
     jsonpath: str
 
 
@@ -110,6 +114,7 @@ class TeardownJsonAssertData(BaseModel):
 
 class TeardownSqlAssertData(TeardownJsonAssertData):
     sql: str
+    fetch: QueryFetchType | None = None
 
 
 class TeardownJsonSchemaAssertData(BaseModel):

diff --git a/httpfpt/utils/assert_control.py b/httpfpt/utils/assert_control.py
@@ -122,13 +122,14 @@ def _sql_asserter(self, assert_text: dict) -> None:
             assert_value = assert_text['value']
             assert_type = assert_text['type']
             assert_sql = assert_text['sql']
+            assert_fetch = assert_text.get('fetch')
             assert_jsonpath = assert_text['jsonpath']
         except KeyError as e:
             raise AssertSyntaxError(f'SQL 断言格式错误, 请检查: {e}')
         else:
             if assert_sql.split(' ')[0].upper() != SqlType.select:
                 raise AssertSyntaxError(f'SQL 断言 {assert_check}:{assert_type} 执行失败，请检查 SQL 是否为 DQL 类型')
-            sql_data = mysql_client.exec_case_sql(assert_sql)
+            sql_data = mysql_client.exec_case_sql(assert_sql, assert_fetch)
             if not isinstance(sql_data, dict):
                 raise JsonPathFindError('jsonpath 取值失败, SQL 语句执行结果不是有效的 dict 类型')
             sql_value = findall(assert_jsonpath, sql_data)