Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/chat/task/llm.py

@@ -536,8 +536,8 @@ def generate_filter(self, sql: str, tables: List):
                     ).first()
                     if obj is not None:
                         res.append(transRecord2DTO(self.session, permission))
-            wheres = transFilterTree(self.session, res, self.ds)
-            filters.append({"table": table.table_name, "filter": wheres})
+            where_str = transFilterTree(self.session, res, self.ds)
+            filters.append({"table": table.table_name, "filter": where_str})
 
         filter = json.dumps(filters, ensure_ascii=False)
 

backend/apps/datasource/crud/datasource.py

@@ -278,8 +278,8 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table
                 ).first()
                 if obj is not None:
                     res.append(transRecord2DTO(session, permission))
-        wheres = transFilterTree(session, res, ds)
-        where = (' where ' + wheres) if wheres is not None and wheres != '' else ''
+        where_str = transFilterTree(session, res, ds)
+        where = (' where ' + where_str) if where_str is not None and where_str != '' else ''
 
     fields = [f.field_name for f in f_list]
     if fields is None or len(fields) == 0:

backend/apps/db/db.py

@@ -254,20 +254,17 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
 
 
 def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str):
-    session = get_session(ds)
-    result = session.execute(text(sql))
-    try:
-        columns = result.keys()._keys
-        res = result.fetchall()
-        result_list = [
-            {columns[i]: float(value) if isinstance(value, Decimal) else value for i, value in enumerate(tuple_item)}
-            for tuple_item in res
-        ]
-        return {"fields": columns, "data": result_list, "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
-    except Exception as ex:
-        raise ex
-    finally:
-        if result is not None:
-            result.close()
-        if session is not None:
-            session.close()
+    with get_session(ds) as session:
+        with session.execute(text(sql)) as result:
+            try:
+                columns = result.keys()._keys
+                res = result.fetchall()
+                result_list = [
+                    {columns[i]: float(value) if isinstance(value, Decimal) else value for i, value in
+                     enumerate(tuple_item)}
+                    for tuple_item in res
+                ]
+                return {"fields": columns, "data": result_list,
+                        "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
+            except Exception as ex:
+                raise ex

backend/apps/system/api/user.py

@@ -1,7 +1,7 @@
 from typing import Optional
 from fastapi import APIRouter, HTTPException, Query
 from sqlmodel import func, or_, select, delete as sqlmodel_delete
-from apps.system.crud.user import check_account_exists, check_email_exists, get_db_user, single_delete, user_ws_options
+from apps.system.crud.user import check_account_exists, check_email_exists, check_email_format, check_pwd_format, get_db_user, single_delete, user_ws_options
 from apps.system.models.system_model import UserWsModel
 from apps.system.models.user import UserModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
@@ -120,6 +120,8 @@ async def create(session: SessionDep, creator: UserCreator):
         raise Exception(f"Account [{creator.account}] already exists!")
     if check_email_exists(session=session, email=creator.email):
         raise Exception(f"Email [{creator.email}] already exists!")
+    if not check_email_format(creator.email):
+        raise Exception(f"Email [{creator.email}] format is invalid!")
     data = creator.model_dump(exclude_unset=True)
     user_model = UserModel.model_validate(data)
     #user_model.create_time = get_timestamp()
@@ -150,6 +152,8 @@ async def update(session: SessionDep, editor: UserEditor):
         raise Exception(f"account cannot be changed!")
     if editor.email != user_model.email and check_email_exists(session=session, account=editor.email):
         raise Exception(f"Email [{editor.email}] already exists!")
+    if not check_email_format(editor.email):
+        raise Exception(f"Email [{editor.email}] format is invalid!")
     origin_oid: int = user_model.oid
     del_stmt = sqlmodel_delete(UserWsModel).where(UserWsModel.uid == editor.id)
     session.exec(del_stmt)
@@ -206,9 +210,12 @@ async def pwdReset(session: SessionDep, current_user: CurrentUser, id: int):
 @router.put("/pwd")
 @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def pwdUpdate(session: SessionDep, current_user: CurrentUser, editor: PwdEditor):
+    new_pwd = editor.new_pwd
+    if not check_pwd_format(new_pwd):
+        raise Exception("Password format is invalid!")
     db_user: UserModel = get_db_user(session=session, user_id=current_user.id)
     if not verify_md5pwd(editor.pwd, db_user.password):
-        raise HTTPException("pwd error")
-    db_user.password = md5pwd(editor.new_pwd)
+        raise Exception(f"pwd [{editor.pwd}] error")
+    db_user.password = md5pwd(new_pwd)
     session.add(db_user)
     session.commit()

backend/apps/system/crud/user.py

@@ -10,6 +10,7 @@
 from common.utils.utils import SQLBotLogUtil
 from ..models.user import UserModel
 from common.core.security import verify_md5pwd
+import re
 
 def get_db_user(*, session: Session, user_id: int) -> UserModel:
     db_user = session.get(UserModel, user_id)
@@ -74,4 +75,23 @@ async def clean_user_cache(id: int):
 def check_account_exists(*, session: Session, account: str) -> bool:
     session.exec(select(func.count()).select_from(UserModel).where(UserModel.account == account)).one() > 0
 def check_email_exists(*, session: Session, email: str) -> bool:
-    return session.exec(select(func.count()).select_from(UserModel).where(UserModel.email == email)).one() > 0
+    return session.exec(select(func.count()).select_from(UserModel).where(UserModel.email == email)).one() > 0
+
+
+# 预编译正则表达式，提高效率
+EMAIL_REGEX = re.compile(
+    r"^[a-zA-Z0-9]+([._-][a-zA-Z0-9]+)*@"
+    r"([a-zA-Z0-9]+(-[a-zA-Z0-9]+)*\.)+"
+    r"[a-zA-Z]{2,}$"
+)
+
+PWD_REGEX = re.compile(
+    r"^(?=.*[a-z])(?=.*[A-Z])(?=.*\d)"
+    r"(?=.*[~!@#$%^&*()_+\-={}|:\"<>?`\[\];',./])"
+    r"[A-Za-z\d~!@#$%^&*()_+\-={}|:\"<>?`\[\];',./]{8,20}$"
+)
+def check_email_format(email: str) -> bool:
+    return bool(EMAIL_REGEX.fullmatch(email))
+
+def check_pwd_format(pwd: str) -> bool:
+    return bool(PWD_REGEX.fullmatch(pwd))