tianfeng_task_modules/config/database.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据库连接配置模块
包含数据库连接参数和SQLAlchemy配置，以及Redis缓存配置
"""

import os
import json
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
import traceback
import sys

class ConfigDict:
    """配置字典类，支持通过点号访问配置项"""
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            if isinstance(value, dict):
                setattr(self, key, ConfigDict(**value))
            else:
                setattr(self, key, value)
    
    def get(self, key, default=None):
        return getattr(self, key, default)
    
    def to_dict(self):
        result = {}
        for key, value in self.__dict__.items():
            if isinstance(value, ConfigDict):
                result[key] = value.to_dict()
            else:
                result[key] = value
        return result

# 数据库连接配置
DB_CONFIG = ConfigDict(
    default=dict(
        dialect='mysql',
        driver='pymysql',
        username='root',
        password='root',
        host='localhost',
        port=3306,
        database='tianfeng_task',
        charset='utf8mb4'
    ),
    test=dict(
        dialect='sqlite',
        database=':memory:'
    )
)

# Redis缓存配置
REDIS_CONFIG = ConfigDict(
    default=dict(
        host='localhost',
        port=6379,
        db=0,
        password=None,
        prefix='tianfeng:',
        socket_timeout=5,
        socket_connect_timeout=5,
        decode_responses=True
    ),
    test=dict(
        host='localhost',
        port=6379,
        db=1,
        password=None,
        prefix='tianfeng_test:',
        decode_responses=True
    )
)

# 当前环境，可通过环境变量设置
ENV = os.environ.get('TIANFENG_ENV', 'default')

# 根据环境获取数据库配置
db_conf = getattr(DB_CONFIG, ENV)

# 构建数据库连接URL
if db_conf.dialect == 'sqlite':
    DATABASE_URL = f"{db_conf.dialect}:///{db_conf.database}"
else:
    DATABASE_URL = (
        f"{db_conf.dialect}+{db_conf.driver}://"
        f"{db_conf.username}:{db_conf.password}@"
        f"{db_conf.host}:{db_conf.port}/{db_conf.database}?"
        f"charset={db_conf.charset}"
    )

# 创建数据库引擎
engine = create_engine(
    DATABASE_URL,
    pool_size=20,
    max_overflow=0,
    pool_recycle=3600,
    pool_pre_ping=True,
    echo=False  # 设置为True可以显示SQL语句，用于调试
)

# 创建会话工厂
SessionFactory = sessionmaker(bind=engine)

# 创建线程安全的会话
db_session = scoped_session(SessionFactory)

# 创建基类
Base = declarative_base()
Base.query = db_session.query_property()

# 数据库配置类
class DBConfig:
    """数据库配置类，提供数据库相关的配置和方法"""
    config = DB_CONFIG
    env = ENV
    url = DATABASE_URL
    engine = engine
    session = db_session
    base = Base
    
    @classmethod
    def get_config(cls):
        """获取当前环境的数据库配置"""
        return getattr(cls.config, cls.env)
    
    @classmethod
    def get_session(cls):
        """获取数据库会话"""
        return cls.session
    
    @classmethod
    def init_db(cls):
        """
        初始化数据库
        创建所有表
        """
        # 测试数据库连接
        try:
            print(f"尝试连接数据库: {cls.url}")
            connection = cls.engine.connect()
            print("数据库连接成功！")
            connection.close()
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            print("详细错误信息:")
            traceback.print_exc(file=sys.stdout)
            raise
        
        # 导入所有模型，确保它们已注册到Base
        import data.models
        
        # 首先尝试创建数据库（如果不存在）
        if cls.get_config().dialect != 'sqlite':
            from sqlalchemy import text
            # 创建一个不指定数据库的连接
            db_conf = cls.get_config()
            try:
                print(f"尝试创建数据库 {db_conf.database} (如果不存在)")
                temp_url = (
                    f"{db_conf.dialect}+{db_conf.driver}://"
                    f"{db_conf.username}:{db_conf.password}@"
                    f"{db_conf.host}:{db_conf.port}/"
                    f"?charset={db_conf.charset}"
                )
                print(f"临时连接URL: {temp_url}")
                temp_engine = create_engine(temp_url)
                with temp_engine.connect() as conn:
                    conn.execute(text(f"CREATE DATABASE IF NOT EXISTS {db_conf.database} CHARACTER SET {db_conf.charset} COLLATE {db_conf.charset}_unicode_ci;"))
                    conn.commit()
                temp_engine.dispose()
                print(f"数据库 {db_conf.database} 创建或已存在")
            except Exception as e:
                print(f"创建数据库失败: {str(e)}")
                print("详细错误信息:")
                traceback.print_exc(file=sys.stdout)
                raise
        
        # 创建所有表
        try:
            print("开始创建所有表...")
            cls.base.metadata.create_all(bind=cls.engine)
            print("所有表创建成功")
        except Exception as e:
            print(f"创建表失败: {str(e)}")
            print("详细错误信息:")
            traceback.print_exc(file=sys.stdout)
            raise
    
    @classmethod
    def shutdown_session(cls, exception=None):
        """
        关闭会话
        在应用程序关闭时调用
        """
        cls.session.remove()

# 缓存配置类
class CacheConfig:
    """缓存配置类，提供Redis缓存相关的配置和方法"""
    config = REDIS_CONFIG
    env = ENV
    _redis_client = None
    
    @classmethod
    def get_config(cls):
        """获取当前环境的Redis配置"""
        return getattr(cls.config, cls.env)
    
    @classmethod
    def get_redis_client(cls):
        """获取Redis客户端实例"""
        if cls._redis_client is None:
            try:
                import redis
                redis_conf = cls.get_config()
                cls._redis_client = redis.Redis(
                    host=redis_conf.host,
                    port=redis_conf.port,
                    db=redis_conf.db,
                    password=redis_conf.password,
                    socket_timeout=getattr(redis_conf, 'socket_timeout', 5),
                    socket_connect_timeout=getattr(redis_conf, 'socket_connect_timeout', 5),
                    decode_responses=getattr(redis_conf, 'decode_responses', True)
                )
            except ImportError:
                raise ImportError("Redis package is not installed. Please install it with 'pip install redis'")
            except Exception as e:
                print(f"Error connecting to Redis: {e}")
                return None
        return cls._redis_client
    
    @classmethod
    def get_key(cls, key):
        """获取带前缀的缓存键"""
        prefix = getattr(cls.get_config(), 'prefix', 'tianfeng:')
        return f"{prefix}{key}"
    
    @classmethod
    def set(cls, key, value, expire=None):
        """
        设置缓存
        
        Args:
            key (str): 缓存键
            value (any): 缓存值，非字符串类型会被JSON序列化
            expire (int, optional): 过期时间（秒）
            
        Returns:
            bool: 是否设置成功
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return False
        
        if not isinstance(value, (str, int, float, bool)):
            value = json.dumps(value)
        
        full_key = cls.get_key(key)
        if expire:
            return redis_client.setex(full_key, expire, value)
        else:
            return redis_client.set(full_key, value)
    
    @classmethod
    def get(cls, key, default=None):
        """
        获取缓存
        
        Args:
            key (str): 缓存键
            default (any, optional): 默认值
            
        Returns:
            any: 缓存值或默认值
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return default
        
        full_key = cls.get_key(key)
        value = redis_client.get(full_key)
        
        if value is None:
            return default
        
        # 尝试解析JSON
        try:
            if value.startswith('{') or value.startswith('['):
                return json.loads(value)
        except (json.JSONDecodeError, AttributeError):
            pass
        
        return value
    
    @classmethod
    def delete(cls, key):
        """
        删除缓存
        
        Args:
            key (str): 缓存键
            
        Returns:
            bool: 是否删除成功
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return False
        
        full_key = cls.get_key(key)
        return redis_client.delete(full_key) > 0
    
    @classmethod
    def exists(cls, key):
        """
        检查缓存是否存在
        
        Args:
            key (str): 缓存键
            
        Returns:
            bool: 是否存在
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return False
        
        full_key = cls.get_key(key)
        return redis_client.exists(full_key) > 0
    
    @classmethod
    def ttl(cls, key):
        """
        获取缓存剩余过期时间
        
        Args:
            key (str): 缓存键
            
        Returns:
            int: 剩余秒数，-1表示永不过期，-2表示不存在
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return -2
        
        full_key = cls.get_key(key)
        return redis_client.ttl(full_key)
    
    @classmethod
    def clear_all(cls):
        """
        清除当前环境下的所有缓存
        
        Returns:
            bool: 是否清除成功
        """
        redis_client = cls.get_redis_client()
        if not redis_client:
            return False
        
        prefix = getattr(cls.get_config(), 'prefix', 'tianfeng:')
        keys = redis_client.keys(f"{prefix}*")
        if keys:
            return redis_client.delete(*keys) > 0
        return True

# 兼容旧代码的函数
def init_db():
    """初始化数据库（兼容旧代码）"""
    DBConfig.init_db()

def shutdown_session(exception=None):
    """关闭会话（兼容旧代码）"""
    DBConfig.shutdown_session(exception)