175 lines
6.0 KiB
Python
Executable File
175 lines
6.0 KiB
Python
Executable File
from functools import wraps
|
|
import time
|
|
from flask import current_app
|
|
import logging
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy import text
|
|
from scripts.models import db
|
|
from collections import defaultdict
|
|
|
|
# 简单的内存缓存实现
|
|
class QueryCache:
|
|
def __init__(self, max_size=100, ttl=300): # 默认缓存300秒
|
|
self.cache = {}
|
|
self.max_size = max_size
|
|
self.ttl = ttl
|
|
self.access_times = defaultdict(int)
|
|
|
|
def get(self, key):
|
|
if key in self.cache:
|
|
item = self.cache[key]
|
|
if time.time() - item['time'] < self.ttl:
|
|
self.access_times[key] += 1
|
|
return item['value']
|
|
else:
|
|
# 缓存过期，删除
|
|
del self.cache[key]
|
|
if key in self.access_times:
|
|
del self.access_times[key]
|
|
return None
|
|
|
|
def set(self, key, value):
|
|
# 如果缓存满了，删除最少访问的项
|
|
if len(self.cache) >= self.max_size:
|
|
# 找出访问次数最少的键
|
|
if self.access_times:
|
|
min_key = min(self.access_times, key=self.access_times.get)
|
|
if min_key in self.cache:
|
|
del self.cache[min_key]
|
|
del self.access_times[min_key]
|
|
|
|
self.cache[key] = {
|
|
'value': value,
|
|
'time': time.time()
|
|
}
|
|
self.access_times[key] = 1
|
|
return value
|
|
|
|
# 创建全局缓存实例
|
|
query_cache = QueryCache()
|
|
|
|
# 数据库操作装饰器 - 用于缓存结果
|
|
def cache_query(ttl=300):
|
|
"""
|
|
缓存查询结果的装饰器
|
|
:param ttl: 缓存有效期（秒）
|
|
"""
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
# 生成缓存键
|
|
cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
|
|
|
|
# 尝试从缓存获取
|
|
cached_result = query_cache.get(cache_key)
|
|
if cached_result is not None:
|
|
return cached_result
|
|
|
|
# 执行查询
|
|
result = func(*args, **kwargs)
|
|
|
|
# 设置缓存
|
|
query_cache.set(cache_key, result)
|
|
return result
|
|
return wrapper
|
|
return decorator
|
|
|
|
# 数据库操作装饰器 - 用于异常处理
|
|
def db_operation(retries=3, retry_delay=0.5):
|
|
"""
|
|
数据库操作的装饰器，提供异常处理和重试功能
|
|
:param retries: 重试次数
|
|
:param retry_delay: 重试延迟(秒)
|
|
"""
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
last_error = None
|
|
for attempt in range(retries):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except SQLAlchemyError as e:
|
|
last_error = e
|
|
current_app.logger.warning(f"数据库操作失败: {str(e)}，尝试重试 ({attempt+1}/{retries})")
|
|
|
|
# 回滚会话
|
|
db.session.rollback()
|
|
|
|
# 如果不是最后一次尝试，则等待后重试
|
|
if attempt < retries - 1:
|
|
time.sleep(retry_delay)
|
|
|
|
# 所有重试都失败
|
|
current_app.logger.error(f"数据库操作失败(已重试{retries}次): {str(last_error)}")
|
|
raise last_error
|
|
return wrapper
|
|
return decorator
|
|
|
|
# 辅助函数：清除特定表相关的所有缓存
|
|
def clear_table_cache(table_name):
|
|
"""
|
|
清除与特定表相关的所有缓存
|
|
:param table_name: 表名
|
|
"""
|
|
keys_to_delete = []
|
|
for key in list(query_cache.cache.keys()):
|
|
if table_name.lower() in key.lower():
|
|
keys_to_delete.append(key)
|
|
|
|
for key in keys_to_delete:
|
|
if key in query_cache.cache:
|
|
del query_cache.cache[key]
|
|
if key in query_cache.access_times:
|
|
del query_cache.access_times[key]
|
|
|
|
# 数据库连接重试机制
|
|
def retry_database_connection(db_instance, app, max_retries=10, retry_delay=5):
|
|
"""
|
|
尝试连接数据库，在连接失败时进行重试
|
|
:param db_instance: SQLAlchemy数据库实例
|
|
:param app: Flask应用实例
|
|
:param max_retries: 最大重试次数
|
|
:param retry_delay: 重试间隔(秒)
|
|
:return: 连接成功返回True，失败返回False
|
|
"""
|
|
import sqlalchemy.exc as sa_exc
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
app.logger.info(f"尝试连接数据库... (尝试 {attempt + 1}/{max_retries})")
|
|
|
|
# 使用with语句确保应用上下文正确
|
|
with app.app_context():
|
|
# 尝试建立连接 - SQLAlchemy 2.0 兼容方式
|
|
with db_instance.engine.connect() as connection:
|
|
connection.execute(text("SELECT 1"))
|
|
connection.commit() # 确保事务提交
|
|
|
|
app.logger.info("数据库连接成功！")
|
|
return True
|
|
|
|
except (sa_exc.OperationalError, sa_exc.DatabaseError, Exception) as e:
|
|
app.logger.warning(f"数据库连接失败: {str(e)}")
|
|
|
|
if attempt < max_retries - 1:
|
|
app.logger.info(f"等待 {retry_delay} 秒后重试...")
|
|
time.sleep(retry_delay)
|
|
else:
|
|
app.logger.error(f"数据库连接失败，已重试 {max_retries} 次，放弃连接")
|
|
return False
|
|
|
|
return False
|
|
|
|
# 通用查询函数，包含异常处理和缓存
|
|
@db_operation()
|
|
@cache_query()
|
|
def get_all_distinct_values(model, column_name):
|
|
"""
|
|
获取指定模型指定列的所有不同值
|
|
:param model: 模型类
|
|
:param column_name: 列名
|
|
:return: 所有不同值的列表
|
|
"""
|
|
column = getattr(model, column_name)
|
|
results = db.session.query(column).distinct().all()
|
|
return [result[0] for result in results if result[0]] |