course-agent-od/backend/scripts/docker_init.py

#!/usr/bin/env python3
"""
Docker环境初始化脚本
"""

import os
import sys
import time
import logging
from pathlib import Path
from datetime import datetime

# 添加src目录到Python路径
sys.path.append(str(Path(__file__).parent.parent))

from src.core.database import SessionLocal, create_tables, check_database_connection, get_database_info
from src.core.config import get_settings, is_postgresql_database
from src.models.user import User
from src.models.knowledge_base import KnowledgeBase
from src.models.forum import ForumCategory
from src.services.auth_service import AuthService

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

class DockerInitializer:
    """Docker环境初始化"""
    
    def __init__(self):
        self.settings = get_settings()
        self.session = SessionLocal()
        self.auth_service = AuthService(self.session)
    
    def initialize(self):
        """执行初始化"""
        logger.info("=" * 60)
        logger.info("Docker环境初始化")
        logger.info("=" * 60)
        
        # 1. 检查数据库连接
        if not self._check_database():
            return False
        
        # 2. 创建数据库表
        self._create_tables()
        
        # 3. 创建默认管理员用户
        self._create_admin_user()
        
        # 4. 创建系统知识库
        self._create_system_knowledge_bases()
        
        # 5. 创建论坛分类
        self._create_forum_categories()
        
        # 6. 检查数据目录
        self._check_data_directories()
        
        logger.info("✅ Docker环境初始化完成")
        return True
    
    def _check_database(self):
        """检查数据库连接"""
        logger.info("检查数据库连接...")
        
        max_retries = 10
        retry_delay = 5
        
        for i in range(max_retries):
            try:
                if check_database_connection():
                    db_info = get_database_info()
                    logger.info(f"✅ 数据库连接成功: {db_info.get('type', 'Unknown')}")
                    
                    if is_postgresql_database():
                        logger.info(f"  数据库: {db_info.get('database', 'Unknown')}")
                        logger.info(f"  版本: {db_info.get('version', 'Unknown')}")
                    
                    return True
                else:
                    logger.warning(f"数据库连接失败，重试 {i+1}/{max_retries}...")
            except Exception as e:
                logger.warning(f"数据库连接异常: {e}")
            
            if i < max_retries - 1:
                time.sleep(retry_delay)
        
        logger.error("❌ 数据库连接失败，请检查数据库服务")
        return False
    
    def _create_tables(self):
        """创建数据库表"""
        logger.info("创建数据库表...")
        
        try:
            create_tables()
            logger.info("✅ 数据库表创建完成")
        except Exception as e:
            logger.error(f"❌ 创建数据库表失败: {e}")
            raise
    
    def _create_admin_user(self):
        """创建默认管理员用户"""
        logger.info("创建默认管理员用户...")
        
        try:
            # 检查是否已存在管理员用户
            admin_user = self.session.query(User).filter(
                User.username == "admin"
            ).first()
            
            if admin_user:
                logger.info("管理员用户已存在")
                return
            
            # 创建管理员用户
            admin_data = {
                "username": "admin",
                "email": "admin@course-agent.local",
                "password": "admin123",  # 默认密码，首次登录后应修改
                "full_name": "系统管理员",
                "is_superuser": True
            }
            
            admin_user = self.auth_service.create_user(**admin_data)
            
            if admin_user:
                logger.info("✅ 管理员用户创建成功")
                logger.info(f"  用户名: {admin_user.username}")
                logger.info(f"  邮箱: {admin_user.email}")
                logger.info("⚠️  请尽快修改默认密码")
            else:
                logger.warning("管理员用户创建失败")
                
        except Exception as e:
            logger.error(f"创建管理员用户失败: {e}")
    
    def _create_system_knowledge_bases(self):
        """创建系统知识库"""
        logger.info("创建系统知识库...")
        
        system_kbs = [
            {
                "name": "国土空间法律法规集",
                "description": "国土空间规划相关法律法规、政策文件",
                "is_system": True,
                "is_active": True
            },
            {
                "name": "国土空间规划原理集",
                "description": "国土空间规划理论、原理、方法",
                "is_system": True,
                "is_active": True
            },
            {
                "name": "国土空间规划案例集",
                "description": "各级国土空间规划案例",
                "is_system": True,
                "is_active": True
            }
        ]
        
        created_count = 0
        for kb_data in system_kbs:
            try:
                # 检查是否已存在
                existing = self.session.query(KnowledgeBase).filter(
                    KnowledgeBase.name == kb_data["name"],
                    KnowledgeBase.is_system == True
                ).first()
                
                if existing:
                    continue
                
                # 创建知识库
                kb = KnowledgeBase(
                    user_id=1,  # 管理员用户ID
                    **kb_data
                )
                
                self.session.add(kb)
                created_count += 1
                
            except Exception as e:
                logger.error(f"创建知识库失败 {kb_data['name']}: {e}")
        
        if created_count > 0:
            self.session.commit()
            logger.info(f"✅ 创建了 {created_count} 个系统知识库")
        else:
            logger.info("系统知识库已存在")
    
    def _create_forum_categories(self):
        """创建论坛分类"""
        logger.info("创建论坛分类...")
        
        categories = [
            {
                "name": "系统使用优化建议",
                "description": "分享系统使用经验，提出改进建议"
            },
            {
                "name": "课程学习反馈",
                "description": "交流学习心得，讨论课程内容"
            }
        ]
        
        created_count = 0
        for cat_data in categories:
            try:
                # 检查是否已存在
                existing = self.session.query(ForumCategory).filter(
                    ForumCategory.name == cat_data["name"]
                ).first()
                
                if existing:
                    continue
                
                # 创建分类
                category = ForumCategory(**cat_data)
                self.session.add(category)
                created_count += 1
                
            except Exception as e:
                logger.error(f"创建论坛分类失败 {cat_data['name']}: {e}")
        
        if created_count > 0:
            self.session.commit()
            logger.info(f"✅ 创建了 {created_count} 个论坛分类")
        else:
            logger.info("论坛分类已存在")
    
    def _check_data_directories(self):
        """检查数据目录"""
        logger.info("检查数据目录...")
        
        directories = [
            self.settings.vector_store_path,
            self.settings.upload_dir,
            self.settings.knowledge_base_dir,
            self.settings.generated_images_dir,
            os.path.dirname(self.settings.log_file),
        ]
        
        for dir_path in directories:
            try:
                path = Path(dir_path)
                path.mkdir(parents=True, exist_ok=True)
                
                # 检查权限
                test_file = path / ".test_write"
                test_file.touch()
                test_file.unlink()
                
                logger.info(f"✅ 目录可访问: {dir_path}")
            except Exception as e:
                logger.error(f"❌ 目录访问失败 {dir_path}: {e}")
    
    def close(self):
        """关闭资源"""
        if self.session:
            self.session.close()

def main():
    """主函数"""
    initializer = None
    
    try:
        initializer = DockerInitializer()
        success = initializer.initialize()
        
        if success:
            print("\n" + "=" * 60)
            print("初始化完成！")
            print("=" * 60)
            print("\n访问信息:")
            print(f"  前端: http://localhost:8001")
            print(f"  后端API: http://localhost:8000")
            print(f"  API文档: http://localhost:8000/docs")
            print("\n默认管理员账号:")
            print("  用户名: admin")
            print("  密码: admin123")
            print("  ⚠️  请尽快修改默认密码")
            print("\n" + "=" * 60)
        else:
            print("❌ 初始化失败")
            sys.exit(1)
            
    except Exception as e:
        print(f"❌ 初始化过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
        
    finally:
        if initializer:
            initializer.close()

if __name__ == "__main__":
    main()