greater_wms/backup/postgresql.py

import os
import subprocess
import logging
import time
from pathlib import Path
import shutil

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('postgres_service_manager.log')
    ]
)
logger = logging.getLogger(__name__)

def get_postgres_service_name():
    """获取 PostgreSQL 服务名称"""
    service_names = [
        "postgresql-x64-16",
        "postgresql-x64-15",
        "postgresql-x64-14",
        "postgresql-x64-13",
        "postgresql-x64-12",
        "postgresql"
    ]
    
    # 检查服务状态
    for service in service_names:
        try:
            result = subprocess.run(
                ["sc", "query", service], 
                capture_output=True, 
                text=True,
                encoding='utf-8',
                errors='ignore'
            )
            if "RUNNING" in result.stdout or "STOPPED" in result.stdout:
                return service
        except Exception:
            continue
    
    return service_names[0] if service_names else "postgresql"

def is_service_running(service_name=None):
    """检查服务是否正在运行"""
    if not service_name:
        service_name = get_postgres_service_name()
    
    try:
        result = subprocess.run(
            ["sc", "query", service_name], 
            capture_output=True, 
            text=True,
            encoding='utf-8',
            errors='ignore'
        )
        return "RUNNING" in result.stdout
    except Exception as e:
        logger.error(f"检查服务状态失败: {str(e)}")
        return False

def stop_postgres_service():
    """停止 PostgreSQL 服务"""
    try:
        service_name = get_postgres_service_name()
        logger.info(f"尝试停止 PostgreSQL 服务: {service_name}")
        
        # 使用 net stop 停止服务
        result = subprocess.run(
            ["net", "stop", service_name], 
            capture_output=True, 
            text=True,
            encoding='utf-8',
            errors='ignore'
        )
        
        if result.returncode == 0:
            logger.info(f"PostgreSQL 服务 '{service_name}' 已停止")
            return True
        
        # 如果 net stop 失败，尝试 sc stop
        result = subprocess.run(
            ["sc", "stop", service_name], 
            capture_output=True, 
            text=True,
            encoding='utf-8',
            errors='ignore'
        )
        
        if result.returncode == 0:
            logger.info(f"PostgreSQL 服务 '{service_name}' 已停止")
            return True
        
        # 如果服务未停止，尝试强制终止进程
        logger.warning("服务未正常停止，尝试终止进程")
        subprocess.run(
            ["taskkill", "/F", "/IM", "postgres.exe"], 
            capture_output=True,
            encoding='utf-8',
            errors='ignore'
        )
        logger.info("PostgreSQL 进程已终止")
        return True
        
    except Exception as e:
        logger.error(f"停止 PostgreSQL 服务失败: {str(e)}")
        return False

def start_postgres_service():
    """启动 PostgreSQL 服务"""
    try:
        service_name = get_postgres_service_name()
        logger.info(f"尝试启动 PostgreSQL 服务: {service_name}")
        
        # 使用 net start 启动服务
        result = subprocess.run(
            ["net", "start", service_name], 
            capture_output=True, 
            text=True,
            encoding='utf-8',
            errors='ignore'
        )
        
        if result.returncode == 0:
            logger.info(f"PostgreSQL 服务 '{service_name}' 已启动")
            return True
        
        # 如果 net start 失败，尝试 sc start
        result = subprocess.run(
            ["sc", "start", service_name], 
            capture_output=True, 
            text=True,
            encoding='utf-8',
            errors='ignore'
        )
        
        if result.returncode == 0:
            logger.info(f"PostgreSQL 服务 '{service_name}' 已启动")
            return True
        
        raise RuntimeError("启动 PostgreSQL 服务失败")
        
    except Exception as e:
        logger.error(f"启动 PostgreSQL 服务失败: {str(e)}")
        return False

def restart_postgres_service():
    """重启 PostgreSQL 服务"""
    try:
        if stop_postgres_service():
            time.sleep(3)  # 等待服务完全停止
        return start_postgres_service()
    except Exception as e:
        logger.error(f"重启 PostgreSQL 服务失败: {str(e)}")
        return False

def get_postgres_data_dir():
    """获取 PostgreSQL 数据目录"""
    possible_paths = [
        r"D:/app/postgresql/data",
        r"C:/Program Files/PostgreSQL/15/data",
        r"C:/Program Files/PostgreSQL/14/data",
        r"C:/Program Files/PostgreSQL/13/data",
        r"C:/Program Files/PostgreSQL/12/data",
    ]
    
    env_path = os.environ.get('PGDATA')
    if env_path and Path(env_path).exists():
        return Path(env_path)
    
    for path in possible_paths:
        if Path(path).exists():
            return Path(path)
    
    raise FileNotFoundError("无法找到 PostgreSQL 数据目录")

def fix_postgres_service():
    """修复 PostgreSQL 服务启动问题"""
    try:
        logger.info("开始修复 PostgreSQL 服务启动问题...")
        
        # 1. 停止服务
        if not stop_postgres_service():
            logger.warning("无法停止服务，尝试强制修复")
        
        # 2. 检查数据目录
        try:
            data_dir = get_postgres_data_dir()
            logger.info(f"找到 PostgreSQL 数据目录: {data_dir}")
            
            # 3. 检查并处理恢复配置文件
            recovery_signal = data_dir / "recovery.signal"
            recovery_done = data_dir / "recovery.done"
            auto_conf = data_dir / "postgresql.auto.conf"
            
            if recovery_signal.exists():
                logger.info("发现 recovery.signal 文件，尝试完成恢复过程")
                if recovery_done.exists():
                    recovery_done.unlink()
                recovery_signal.rename(recovery_done)
                logger.info("已将 recovery.signal 重命名为 recovery.done")
            
            # 4. 删除可能引起问题的自动配置文件
            if auto_conf.exists():
                backup_auto_conf = auto_conf.with_suffix(".auto.conf.bak")
                shutil.copy(auto_conf, backup_auto_conf)
                auto_conf.unlink()
                logger.info("已备份并删除 postgresql.auto.conf 文件")
            
            # 5. 检查 postgresql.conf 文件
            conf_file = data_dir / "postgresql.conf"
            if conf_file.exists():
                logger.info("检查 postgresql.conf 文件完整性")
                # 这里可以添加更多的配置文件检查逻辑
            
        except FileNotFoundError as e:
            logger.warning(f"无法找到数据目录: {str(e)}")
        
        # 6. 确保所有 PostgreSQL 进程已终止
        subprocess.run(
            ["taskkill", "/F", "/IM", "postgres.exe"], 
            capture_output=True,
            encoding='utf-8',
            errors='ignore'
        )
        
        # 7. 尝试启动服务
        time.sleep(2)
        if start_postgres_service():
            logger.info("PostgreSQL 服务修复成功")
            return True
        else:
            logger.error("修复后仍无法启动服务")
            return False
            
    except Exception as e:
        logger.error(f"修复过程中发生错误: {str(e)}")
        return False

def check_service_status():
    """检查服务状态并显示详细信息"""
    service_name = get_postgres_service_name()
    is_running = is_service_running(service_name)
    
    print(f"服务名称: {service_name}")
    print(f"运行状态: {'运行中' if is_running else '已停止'}")
    
    if is_running:
        try:
            # 尝试连接数据库验证服务状态
            import psycopg2
            try:
                conn = psycopg2.connect(
                    dbname='postgres',
                    user='postgres',
                    password='abc@1234',
                    host='localhost',
                    port='5432'
                )
                conn.close()
                print("数据库连接: 正常")
            except Exception as e:
                print(f"数据库连接: 异常 ({str(e)})")
        except ImportError:
            print("数据库连接: 未安装 psycopg2 库，无法测试连接")
    
    return is_running

def main():
    """主函数：提供命令行界面"""
    print("=" * 50)
    print("PostgreSQL 服务管理工具")
    print("=" * 50)
    
    while True:
        print("\n请选择操作:")
        print("1. 检查服务状态")
        print("2. 启动服务")
        print("3. 停止服务")
        print("4. 重启服务")
        print("5. 修复服务启动问题")
        print("6. 退出")
        
        choice = input("请输入选项 (1-6): ").strip()
        
        if choice == "1":
            print("\n检查服务状态...")
            check_service_status()
            
        elif choice == "2":
            print("\n启动服务...")
            if start_postgres_service():
                print("服务启动成功")
            else:
                print("服务启动失败，请尝试修复")
                
        elif choice == "3":
            print("\n停止服务...")
            if stop_postgres_service():
                print("服务停止成功")
            else:
                print("服务停止失败")
                
        elif choice == "4":
            print("\n重启服务...")
            if restart_postgres_service():
                print("服务重启成功")
            else:
                print("服务重启失败")
                
        elif choice == "5":
            print("\n修复服务启动问题...")
            if fix_postgres_service():
                print("服务修复成功")
            else:
                print("服务修复失败，请查看日志文件")
                
        elif choice == "6":
            print("退出程序")
            break
            
        else:
            print("无效选项，请重新选择")
            
        # input("\n按回车键继续...")

if __name__ == "__main__":
    # 直接运行测试
    print("正在启动 PostgreSQL 服务管理工具...")
    main()