139 lines
4.9 KiB
Python
139 lines
4.9 KiB
Python
"""数据库连接池和查询函数"""
|
||
import pymysql
|
||
from dbutils.pooled_db import PooledDB
|
||
from config import DB_CONFIG
|
||
|
||
_pool = None
|
||
|
||
def get_pool():
|
||
global _pool
|
||
if _pool is None:
|
||
_pool = PooledDB(
|
||
creator=pymysql,
|
||
maxconnections=10,
|
||
mincached=2,
|
||
blocking=True,
|
||
host=DB_CONFIG["host"],
|
||
port=DB_CONFIG["port"],
|
||
user=DB_CONFIG["user"],
|
||
password=DB_CONFIG["password"],
|
||
database=DB_CONFIG["database"],
|
||
charset="utf8mb4",
|
||
cursorclass=pymysql.cursors.DictCursor
|
||
)
|
||
return _pool
|
||
|
||
def get_db():
|
||
return get_pool().connection()
|
||
|
||
def _client_ip():
|
||
"""优先取 X-Forwarded-For 首个地址，其次 remote_addr（与审计模块一致）。"""
|
||
from flask import request
|
||
xff = request.headers.get('X-Forwarded-For', '')
|
||
if xff:
|
||
first = xff.split(',')[0].strip()
|
||
if first:
|
||
return first
|
||
return request.remote_addr or ''
|
||
|
||
|
||
def add_operation_log(action, detail, user_id=None, ip_address='', page='', field='',
|
||
event_type='', result_status='success', risk_level='', object_type='',
|
||
object_id='', old_values=None, new_values=None):
|
||
"""记录可读业务日志，同时写入结构化审计字段。"""
|
||
import json
|
||
from flask import session, g, request
|
||
from services.audit_schema import (ensure_audit_schema, infer_page, classify_event,
|
||
classify_risk, infer_object)
|
||
ensure_audit_schema()
|
||
uid = user_id or session.get('user_id')
|
||
ip = ip_address or _client_ip()
|
||
path = request.path if request else ''
|
||
method = request.method if request else ''
|
||
submitted = request.get_json(silent=True) if request else None
|
||
inferred_type, inferred_id = infer_object(submitted, path)
|
||
if uid and not session.get('position'):
|
||
try:
|
||
rows = query("SELECT position FROM dashboard_users WHERE id=%s LIMIT 1", (uid,))
|
||
position = rows[0].get('position') if rows else ''
|
||
except Exception:
|
||
position = ''
|
||
else:
|
||
position = session.get('position', '')
|
||
def dump(value):
|
||
return json.dumps(value, ensure_ascii=False, default=str) if value is not None else None
|
||
log_id = query(
|
||
"""INSERT INTO dashboard_logs
|
||
(user_id,action,detail,ip_address,page,field,event_type,result_status,risk_level,
|
||
object_type,object_id,old_values,new_values,request_method,request_path,request_id,
|
||
user_agent,role_snapshot,position_snapshot,http_status)
|
||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
|
||
(uid, action, detail, ip, page or infer_page(path, action) or None, field or None,
|
||
event_type or classify_event(action, path, method), result_status,
|
||
risk_level or classify_risk(action, field, path), object_type or inferred_type or None,
|
||
object_id or inferred_id or None, dump(old_values), dump(new_values if new_values is not None else submitted),
|
||
method or None, path or None, getattr(g, 'audit_request_id', None),
|
||
(request.headers.get('User-Agent','')[:500] if request else None), session.get('role'),
|
||
position or None, None)
|
||
)
|
||
try:
|
||
if not hasattr(g, 'audit_log_ids'): g.audit_log_ids = []
|
||
if log_id: g.audit_log_ids.append(log_id)
|
||
except Exception: pass
|
||
try:
|
||
from services.audit_governance import finalize_log
|
||
finalize_log(log_id)
|
||
except Exception:
|
||
pass
|
||
return log_id
|
||
|
||
|
||
def query(sql, params=None):
|
||
conn = get_db()
|
||
try:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, params or ())
|
||
if sql.strip().upper().startswith("SELECT"):
|
||
return cur.fetchall()
|
||
conn.commit()
|
||
return cur.lastrowid
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def query_iter(sql, params=None, batch_size=2000):
|
||
"""流式查询 SELECT：按批取行，避免一次性把全表读进内存。
|
||
|
||
使用服务端游标（SSDictCursor），配合导出（openpyxl write_only / csv）逐批消费。
|
||
"""
|
||
conn = get_db()
|
||
try:
|
||
cur = conn.cursor(pymysql.cursors.SSDictCursor)
|
||
try:
|
||
cur.execute(sql, params or ())
|
||
while True:
|
||
rows = cur.fetchmany(batch_size)
|
||
if not rows:
|
||
break
|
||
for r in rows:
|
||
yield r
|
||
finally:
|
||
cur.close()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def execute(sql, params=None):
|
||
"""执行写语句并返回受影响行数（rowcount）。
|
||
|
||
用于需要原子抢占/幂等控制的场景（如确认下单防重复）。
|
||
"""
|
||
conn = get_db()
|
||
try:
|
||
with conn.cursor() as cur:
|
||
affected = cur.execute(sql, params or ())
|
||
conn.commit()
|
||
return affected
|
||
finally:
|
||
conn.close()
|