AutoInfo/base/base_daily_logs_generate.py

# -*- coding: utf-8 -*-
"""
用于青龙面板的日志初始化脚本
每天 00:00:00 执行，创建当天的日志记录
"""
import os
import sys
import time
import logging
from datetime import datetime
from typing import Dict, Any

# 添加项目路径
sys.path.append(os.path.join(os.path.abspath(__file__).split('AutoInfo')[0] + 'AutoInfo'))

from utils.utils import *

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


class Config:
    """配置管理类"""
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._load_config()
        return cls._instance

    def _load_config(self):
        """加载配置"""
        try:
            config_json = LoadConfig().load_config()
            self.PROJECT_NAME = config_json.get('PROJECT_NAME', 'AutoInfo')
            self.DB_USER = config_json.get('DB_USER', '')
            self.DB_PASSWORD = config_json.get('DB_PASSWORD', '')
            self.DB_IP = config_json.get('DB_IP', 'localhost')
            self.DB_PORT = config_json.get('DB_PORT', 27017)
            self.MAIL_HOST = config_json.get('MAIL_HOST', '')
            self.MAIL_USER = config_json.get('MAIL_USER', '')
            self.MAIL_PASS = config_json.get('MAIL_PASS', '')
            self.MAIL_SENDER = config_json.get('MAIL_SENDER', '')
            self.MAIL_RECEIVERS = config_json.get('MAIL_RECEIVERS', [])

            # 构建MongoDB连接字符串
            if self.DB_USER and self.DB_PASSWORD:
                self.MONGO_LINK = f'mongodb://{self.DB_USER}:{self.DB_PASSWORD}@{self.DB_IP}:{self.DB_PORT}/?authSource=admin'
            else:
                self.MONGO_LINK = f'mongodb://{self.DB_IP}:{self.DB_PORT}/'

        except Exception as e:
            logger.error(f"加载配置失败: {e}")
            raise


class MongoHandle:
    """MongoDB操作处理类"""

    def __init__(self, db: str, collection: str, del_db: bool = False,
                 del_collection: bool = False, auto_remove: int = 0):
        """
        初始化MongoDB连接

        Args:
            db: 数据库名
            collection: 集合名
            del_db: 是否删除数据库
            del_collection: 是否删除集合
            auto_remove: 自动删除数据的天数阈值
        """
        self.config = Config()
        logger.info(f"连接数据库: {self.config.MONGO_LINK}")

        try:
            self.client = pymongo.MongoClient(
                self.config.MONGO_LINK,
                serverSelectionTimeoutMS=5000  # 5秒超时
            )
            # 测试连接
            self.client.admin.command('ismaster')

            self.db_name = db
            self.collection_name = collection

            self._setup_database(del_db, del_collection)

            if auto_remove > 0:
                self.auto_remove_data(auto_remove)

        except pymongo.errors.ServerSelectionTimeoutError:
            logger.error("无法连接到MongoDB服务器")
            raise
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            raise

    def _setup_database(self, del_db: bool, del_collection: bool):
        """设置数据库和集合"""
        if del_db and self.db_name:
            if self.db_name in self.client.list_database_names():
                self.client.drop_database(self.db_name)
                logger.info(f"已删除数据库: {self.db_name}")

        self.db = self.client[self.db_name]

        if del_collection and self.collection_name:
            if self.collection_name in self.db.list_collection_names():
                self.db.drop_collection(self.collection_name)
                logger.info(f"已删除集合: {self.collection_name}")

        self.collection = self.db[self.collection_name]

    def write_data(self, data: Dict[str, Any]) -> bool:
        """写入数据"""
        try:
            result = self.collection.insert_one(data)
            logger.debug(f"数据插入成功, ID: {result.inserted_id}")
            return True
        except Exception as e:
            logger.error(f"数据插入失败: {e}")
            return False

    def auto_remove_data(self, days: int):
        """自动删除指定天数前的数据"""
        try:
            cutoff_time = int(time.time()) - days * 24 * 60 * 60
            result = self.collection.delete_many({
                'create_time': {'$lt': cutoff_time}
            })
            if result.deleted_count > 0:
                logger.info(f"已删除 {result.deleted_count} 条过期数据")
        except Exception as e:
            logger.error(f"自动删除数据失败: {e}")

    def close(self):
        """关闭数据库连接"""
        if hasattr(self, 'client'):
            self.client.close()
            logger.info("数据库连接已关闭")


class LogsHandler:
    """日志处理类"""

    def __init__(self):
        self.config = Config()
        self.now_day = datetime.now().strftime('%Y-%m-%d')
        self.mongo = None
        self._setup_mongo()

    def _setup_mongo(self):
        """设置MongoDB连接"""
        try:
            self.mongo = MongoHandle(
                db='logs',
                collection=f'logs_{self.now_day}',
                del_db=False,
                del_collection=False,
                auto_remove=0  # 不自动删除，由其他机制处理
            )
        except Exception as e:
            logger.error(f"初始化MongoDB连接失败: {e}")
            raise

    def logs_generate(self) -> bool:
        """生成当天的日志记录"""
        data_to_insert = {
            "title": "daily_log_init",
            "context": f"Daily log collection created for {self.now_day}",
            "state": "created",
            "create_time": int(time.time()),
            "create_datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "project": self.config.PROJECT_NAME
        }

        try:
            success = self.mongo.write_data(data_to_insert)
            if success:
                logger.info(f"成功创建当天日志记录: {self.now_day}")
            return success
        except Exception as e:
            logger.error(f"创建日志记录失败: {e}")
            return False

    def cleanup(self):
        """清理资源"""
        if self.mongo:
            self.mongo.close()


def main():
    """主函数"""
    logger.info("开始创建当天日志记录...")

    logs_handler = None
    try:
        logs_handler = LogsHandler()
        success = logs_handler.logs_generate()

        if success:
            logger.info("当天日志记录创建成功")
            return 0
        else:
            logger.error("当天日志记录创建失败")
            return 1

    except Exception as e:
        logger.error(f"执行过程中发生错误: {e}")
        return 1
    finally:
        if logs_handler:
            logs_handler.cleanup()


if __name__ == '__main__':
    exit_code = main()
    sys.exit(exit_code)