AcidMap/scripts/db_health_check.py

"""
数据库健康检查工具
@description: 检查数据库连接状态、迁移状态和表结构完整性
@author: AcidMap Team
@version: 1.0.0
"""

import os
import sys
import logging
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.database import engine, Base
from alembic.config import Config
from alembic.runtime.migration import MigrationContext
from alembic.script import ScriptDirectory
from sqlalchemy import text
import traceback

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class DatabaseHealthChecker:
    """
    数据库健康检查器
    
    @description: 全面检查数据库状态，包括连接、迁移、表结构等
    """
    
    def __init__(self):
        self.engine = engine
        self.alembic_cfg = Config(os.path.join(os.path.dirname(os.path.dirname(__file__)), "alembic.ini"))
        self.results = {
            "timestamp": datetime.now().isoformat(),
            "connection": False,
            "migration_status": "unknown",
            "tables_exist": False,
            "issues": [],
            "recommendations": []
        }
    
    def check_database_connection(self):
        """
        检查数据库连接
        
        @returns: bool 连接是否成功
        """
        try:
            with self.engine.connect() as conn:
                # 执行简单查询测试连接
                result = conn.execute(text("SELECT 1"))
                result.fetchone()
                
            self.results["connection"] = True
            logger.info("✓ 数据库连接正常")
            return True
            
        except Exception as e:
            self.results["connection"] = False
            self.results["issues"].append(f"数据库连接失败: {str(e)}")
            logger.error(f"✗ 数据库连接失败: {str(e)}")
            return False
    
    def check_migration_status(self):
        """
        检查数据库迁移状态
        
        @returns: dict 迁移状态信息
        """
        try:
            with self.engine.connect() as connection:
                context = MigrationContext.configure(connection)
                current_rev = context.get_current_revision()
                
            script_dir = ScriptDirectory.from_config(self.alembic_cfg)
            head_rev = script_dir.get_current_head()
            
            migration_info = {
                "current_revision": current_rev,
                "head_revision": head_rev,
                "is_up_to_date": current_rev == head_rev
            }
            
            if migration_info["is_up_to_date"]:
                self.results["migration_status"] = "up_to_date"
                logger.info(f"✓ 数据库迁移状态正常 (版本: {current_rev})")
            else:
                self.results["migration_status"] = "outdated"
                self.results["issues"].append(f"数据库版本过时: 当前 {current_rev}, 最新 {head_rev}")
                self.results["recommendations"].append("执行 'python db_migrate.py upgrade' 升级数据库")
                logger.warning(f"⚠ 数据库需要升级: {current_rev} -> {head_rev}")
            
            return migration_info
            
        except Exception as e:
            self.results["migration_status"] = "error"
            self.results["issues"].append(f"迁移状态检查失败: {str(e)}")
            logger.error(f"✗ 迁移状态检查失败: {str(e)}")
            return {"error": str(e)}
    
    def check_table_structure(self):
        """
        检查表结构完整性
        
        @returns: dict 表结构检查结果
        """
        try:
            with self.engine.connect() as conn:
                # 获取所有表名
                tables_query = text("""
                    SELECT table_name 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public'
                """)
                existing_tables = [row[0] for row in conn.execute(tables_query).fetchall()]
                
                # 获取模型定义的表名
                model_tables = [table.name for table in Base.metadata.tables.values()]
                
                # 检查缺失的表
                missing_tables = set(model_tables) - set(existing_tables)
                extra_tables = set(existing_tables) - set(model_tables) - {'alembic_version'}
                
                table_info = {
                    "existing_tables": existing_tables,
                    "model_tables": model_tables,
                    "missing_tables": list(missing_tables),
                    "extra_tables": list(extra_tables)
                }
                
                if missing_tables:
                    self.results["tables_exist"] = False
                    self.results["issues"].append(f"缺失表: {', '.join(missing_tables)}")
                    self.results["recommendations"].append("执行数据库升级或重新创建表结构")
                    logger.warning(f"⚠ 缺失表: {', '.join(missing_tables)}")
                else:
                    self.results["tables_exist"] = True
                    logger.info("✓ 所有必需的表都存在")
                
                if extra_tables:
                    logger.info(f"发现额外表: {', '.join(extra_tables)}")
                
                return table_info
                
        except Exception as e:
            self.results["tables_exist"] = False
            self.results["issues"].append(f"表结构检查失败: {str(e)}")
            logger.error(f"✗ 表结构检查失败: {str(e)}")
            return {"error": str(e)}
    
    def check_spatial_extensions(self):
        """
        检查空间扩展 (PostGIS)
        
        @returns: dict 空间扩展状态
        """
        try:
            with self.engine.connect() as conn:
                # 检查PostGIS扩展
                postgis_query = text("""
                    SELECT extname, extversion 
                    FROM pg_extension 
                    WHERE extname = 'postgis'
                """)
                postgis_result = conn.execute(postgis_query).fetchall()
                
                if postgis_result:
                    version = postgis_result[0][1]
                    logger.info(f"✓ PostGIS 扩展已安装 (版本: {version})")
                    return {"installed": True, "version": version}
                else:
                    self.results["issues"].append("PostGIS 扩展未安装")
                    self.results["recommendations"].append("安装 PostGIS 扩展: CREATE EXTENSION postgis;")
                    logger.warning("⚠ PostGIS 扩展未安装")
                    return {"installed": False}
                    
        except Exception as e:
            self.results["issues"].append(f"空间扩展检查失败: {str(e)}")
            logger.error(f"✗ 空间扩展检查失败: {str(e)}")
            return {"error": str(e)}
    
    def run_full_check(self):
        """
        执行完整的健康检查
        
        @returns: dict 完整的检查结果
        """
        logger.info("开始数据库健康检查...")
        logger.info("=" * 50)
        
        # 检查数据库连接
        if not self.check_database_connection():
            logger.error("数据库连接失败，跳过后续检查")
            return self.results
        
        # 检查迁移状态
        migration_info = self.check_migration_status()
        
        # 检查表结构
        table_info = self.check_table_structure()
        
        # 检查空间扩展
        spatial_info = self.check_spatial_extensions()
        
        # 汇总结果
        self.results.update({
            "migration_info": migration_info,
            "table_info": table_info,
            "spatial_info": spatial_info
        })
        
        logger.info("=" * 50)
        logger.info("健康检查完成")
        
        # 打印总结
        if self.results["issues"]:
            logger.warning(f"发现 {len(self.results['issues'])} 个问题:")
            for issue in self.results["issues"]:
                logger.warning(f"  - {issue}")
        
        if self.results["recommendations"]:
            logger.info(f"建议执行 {len(self.results['recommendations'])} 项操作:")
            for rec in self.results["recommendations"]:
                logger.info(f"  - {rec}")
        
        if not self.results["issues"]:
            logger.info("✓ 数据库状态良好")
        
        return self.results

def main():
    """
    主函数
    """
    try:
        checker = DatabaseHealthChecker()
        results = checker.run_full_check()
        
        # 根据结果设置退出码
        if results["issues"]:
            sys.exit(1)  # 有问题时退出码为1
        else:
            sys.exit(0)  # 正常时退出码为0
            
    except Exception as e:
        logger.error(f"健康检查执行失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        sys.exit(2)  # 检查失败时退出码为2

if __name__ == "__main__":
    main()