Upload app.py

app.py

@@ -28,17 +28,26 @@ def get_db_connection():
 
 
 def save_to_db(table, data):
-    conn = None  # 确保 conn 在函数开始时就被定义
+    conn = None
     try:
+        # 表名白名单验证
+        valid_tables = ["entities", "relations"]
+        if table not in valid_tables:
+            raise ValueError(f"Invalid table: {table}")
+
         conn = get_db_connection()
         with conn.cursor() as cursor:
-            placeholders = ', '.join(['%s'] * len(data))
+            # 使用参数化查询避免注入
             columns = ', '.join(data.keys())
+            placeholders = ', '.join(['%s'] * len(data))
             sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
             cursor.execute(sql, list(data.values()))
         conn.commit()
-    except Exception as e:
-        print(f"数据库写入失败: {e}")
+    except pymysql.Error as e:  # 细化异常类型
+        print(f"数据库错误: {e}")
+        conn.rollback()
+    except ValueError as e:  # 表名无效
+        print(f"参数错误: {e}")
     finally:
         if conn:
             conn.close()