aiData/DbKit/SqlTemplateLoader.py

import logging
import os
import re
import asyncio
import aiofiles
import json

# 导入RedisKit
from Util.RedisKit import redisKit

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

# Redis键名常量
REDIS_SQL_TEMPLATES_KEY = "sql_templates:all_templates"
REDIS_SQL_LOADED_FILES_KEY = "sql_templates:loaded_files"
REDIS_SQL_TEMPLATE_MAP_KEY = "sql_templates:template_map"
REDIS_SQL_TEMPLATES_LOADED_KEY = "sql_templates:templates_loaded"


class SqlTemplateLoader:
    """SQL模板加载器，负责加载和解析SQL模板文件（异步版本）"""

    def __init__(self, sql_dir=None, auto_load=True):
        """初始化SQL模板加载器

        Args:
            sql_dir: SQL模板文件所在目录
            auto_load: 是否自动加载目录下所有SQL模板文件
        """
        self.sql_dir = sql_dir
        self.templates = {}  # 使用嵌套字典存储命名空间下的SQL模板
        self._loaded_files = set()  # 用于记录已加载的文件，避免重复加载
        self._template_map = {}  # 使用平面字典存储完整的模板名称映射，用于快速查找

        # 注意：异步版本不会在__init__中自动加载，需要单独调用异步方法

    async def load_all_templates(self):
        """自动加载SQL目录下所有的SQL模板文件（异步版本）

        Returns:
            int: 成功加载的模板文件数量
        """
        logger.info(f"开始加载SQL模板目录: {self.sql_dir}")

        if not os.path.exists(self.sql_dir):
            logger.warning(f"SQL模板目录不存在: {self.sql_dir}")
            return 0
            
        # 检查是否已经从Redis加载过
        templates_loaded = await self._check_templates_loaded_from_redis()
        if templates_loaded:
            await self._load_from_redis()
            logger.info(f"从Redis加载SQL模板完成，共{len(self._loaded_files)}个文件，{len(self._template_map)}个模板")
            return len(self._loaded_files)

        loaded_count = 0
        sql_files = []

        try:
            # 遍历目录下所有.sql文件
            for filename in os.listdir(self.sql_dir):
                if filename.lower().endswith('.sql'):
                    file_path = os.path.join(self.sql_dir, filename)
                    file_abs_path = os.path.abspath(file_path)
                    sql_files.append(file_abs_path)

            # 使用异步方式加载所有模板，不立即构建映射以提高性能
            for file_abs_path in sql_files:
                # 跳过已经加载过的文件
                if file_abs_path in self._loaded_files:
                    logger.debug(f"SQL文件已加载，跳过: {file_abs_path}")
                    continue

                try:
                    filename = os.path.basename(file_abs_path)
                    if await self.load_template(file_abs_path, build_map=False):
                        loaded_count += 1
                except Exception as e:
                    logger.error(f"加载SQL模板文件 {filename} 失败: {str(e)}")

            # 所有模板加载完成后，统一构建映射
            if loaded_count > 0:
                self._build_template_map()
                logger.info(f"成功加载 {loaded_count} 个SQL模板文件，已构建模板映射")

            # 加载完成后保存到Redis
            await self._save_to_redis()
            return loaded_count
        except Exception as e:
            logger.error(f"遍历SQL模板目录时出错: {str(e)}")
            return 0

    async def load_template(self, template_path, build_map=True):
        """加载指定的SQL模板文件（异步版本）

        Args:
            template_path: SQL模板文件的路径，可以是绝对路径或相对于sql_dir的相对路径
            build_map: 是否立即构建模板映射，默认为True。在批量加载时可设为False以提高性能

        Returns:
            bool: 加载是否成功

        Raises:
            FileNotFoundError: 当模板文件不存在时
            Exception: 当加载过程发生其他错误时
        """
        # 确定模板文件的完整路径
        if self.sql_dir and not os.path.isabs(template_path):
            file_path = os.path.join(self.sql_dir, template_path)
        else:
            file_path = template_path

        # 检查文件是否已加载，避免重复加载
        file_abs_path = os.path.abspath(file_path)
        if file_abs_path in self._loaded_files:
            logger.debug(f"SQL模板文件已加载，跳过: {file_abs_path}")
            return True

        try:
            # 异步读取文件内容
            async with aiofiles.open(file_path, 'r', encoding='utf-8') as f:
                content = await f.read()

            # 使用栈结构解析整个内容
            templates = self._parse_with_stack(content)

            # 处理命名空间
            if 'namespace' in templates and templates['namespace']:
                namespace = templates['namespace']
                sql_templates = {k: v for k, v in templates.items() if k != 'namespace'}
            else:
                # 没有命名空间的情况
                namespace = 'default'
                sql_templates = templates

            # 将解析出的SQL模板添加到templates字典中
            if namespace not in self.templates:
                self.templates[namespace] = {}

            self.templates[namespace].update(sql_templates)

            # 重新构建模板映射，确保新加载的模板可以被快速查找
            if build_map:
                self._build_template_map()

            # 记录已加载的文件
            self._loaded_files.add(file_abs_path)

            return True
        except FileNotFoundError:
            logger.error(f"SQL模板文件不存在: {file_path}")
            raise FileNotFoundError(f"SQL模板文件不存在: {file_path}")
        except Exception as e:
            logger.error(f"加载SQL模板失败: {str(e)}", exc_info=True)
            raise

    def _parse_with_stack(self, content):
        """使用栈结构解析SQL模板，支持嵌套的标签匹配
        支持--单行注释和/* */多行注释

        Args:
            content: 要解析的SQL模板内容

        Returns:
            解析后的SQL模板字典，键为模板名称，值为模板内容
        """
        templates = {}
        stack = []  # 用于存储标签信息的栈
        current_content = []  # 当前正在收集的内容
        in_multiline_comment = False  # 多行注释标记

        try:
            # 按行处理内容
            lines = content.split('\n')
            for line_num, line in enumerate(lines):
                # 处理多行注释
                if in_multiline_comment:
                    # 检查是否是多行注释结束
                    if '*/' in line:
                        in_multiline_comment = False
                        # 截取注释结束后的内容
                        line = line.split('*/', 1)[1]
                    else:
                        # 完全在多行注释内，跳过此行
                        continue

                # 处理单行注释和检查多行注释开始
                if not line.strip().startswith('#') and stack and stack[-1][0] == 'sql':
                    # 检查是否有--注释
                    if '--' in line:
                        # 保留--前面的内容
                        line = line.split('--', 1)[0]

                    # 检查是否有/*多行注释开始
                    if '/*' in line:
                        parts = line.split('/*', 1)
                        # 如果/*和*/在同一行
                        if '*/' in parts[1]:
                            # 只保留/*前面和*/后面的内容
                            comment_part = parts[1].split('*/', 1)
                            line = parts[0] + comment_part[1]
                        else:
                            # 多行注释开始，只保留/*前面的内容
                            line = parts[0]
                            in_multiline_comment = True

                stripped_line = line.strip()

                # 检查是否是标签行
                if stripped_line.startswith('#'):
                    # 检查是否是命名空间标签
                    if stripped_line.startswith('#namespace(') and (
                            stripped_line.find('"') > 0 or stripped_line.find("'") > 0):
                        # 提取命名空间名称
                        quote_pattern = re.search(r'#namespace\((["\'])([^"\']+)\1\)', stripped_line)
                        if quote_pattern:
                            namespace = quote_pattern.group(2)
                            templates['namespace'] = namespace
                            # 将命名空间标签入栈
                            stack.append(('namespace', namespace, line_num))

                    # 检查是否是SQL模板开始标签
                    elif stripped_line.startswith('#sql(') and (
                            stripped_line.find('"') > 0 or stripped_line.find("'") > 0):
                        # 提取SQL模板名称
                        quote_pattern = re.search(r'#sql\((["\'])([^"\']+)\1\)', stripped_line)
                        if quote_pattern:
                            sql_name = quote_pattern.group(2)
                            # 将SQL模板开始标签入栈
                            stack.append(('sql', sql_name, line_num))
                            current_content = []  # 重置内容收集器

                    # 检查是否是#if条件开始标签
                    elif stripped_line.startswith('#if(') and ')' in stripped_line:
                        # 提取条件表达式
                        condition_pattern = re.search(r'#if\(([^)]+)\)', stripped_line)
                        if condition_pattern:
                            condition = condition_pattern.group(1)
                            stack.append(('if', condition, line_num))
                            # 保留#if行到内容中 - 只要有SQL标签在栈中就保留
                            if any(tag[0] == 'sql' for tag in stack):
                                current_content.append(line.rstrip())

                    # 检查是否是结束标签
                    elif stripped_line == '#end' and stack:
                        # 获取栈顶元素
                        tag_type, tag_name, start_line = stack[-1]

                        # 只有当标签不是SQL结束标签时，才保留#end行到内容中
                        if any(tag[0] == 'sql' for tag in stack) and tag_type != 'sql':
                            current_content.append(line.rstrip())

                        # 弹出当前标签
                        stack.pop()

                        # 如果是SQL模板结束标签，保存内容
                        if tag_type == 'sql':
                            sql_content = '\n'.join(current_content).strip()
                            if sql_content:
                                templates[tag_name] = sql_content
                                # 重置内容收集器，准备下一个模板
                                current_content = []

                # 收集内容（不是标签行且当前有活动的SQL模板）
                else:
                    # 只要栈中有SQL标签，就收集内容（不管是在条件块内还是外）
                    if any(tag[0] == 'sql' for tag in stack) and line.strip():
                        # 保留原始缩进
                        current_content.append(line.rstrip())

            # 检查是否所有标签都正确闭合
            if stack:
                unclosed_tags = [f"{tag_type}('{tag_name}')" for tag_type, tag_name, _ in stack]
                logger.warning(f"存在未闭合的标签: {', '.join(unclosed_tags)}")
            else:
                logger.debug("所有标签都已正确闭合")

        except Exception as e:
            logger.error(f"使用栈结构解析SQL模板时出错: {str(e)}", exc_info=True)

        return templates

    async def get_all_template_names(self):
        """
获取所有已加载的SQL模板的完整名称列表（异步版本）

        Returns:
            list: 包含所有模板名称的列表，格式为 "命名空间.模板名称"
        """
        template_names = []
        
        # 遍历所有命名空间
        for namespace, templates in self.templates.items():
            # 遍历当前命名空间下的所有模板
            for template_name in templates.keys():
                # 构建完整的模板名称（命名空间.模板名称）
                full_name = f"{namespace}.{template_name}"
                template_names.append(full_name)
        
        return template_names
        
    def _has_template(self, full_template_name):
        """
        内部方法：快速检查指定的完整模板名称是否存在
        使用MAP机制实现O(1)时间复杂度的查找
        
        Args:
            full_template_name: 完整的模板名称，格式为 "命名空间.模板名称"
            
        Returns:
            bool: 如果模板存在返回True，否则返回False
        """
        logger.debug(f"检查模板是否存在: {full_template_name}")
        
        # 如果_map为空，先构建映射
        if not self._template_map and self.templates:
            logger.debug("模板映射为空，正在构建映射...")
            self._build_template_map()
            logger.debug(f"映射构建完成，包含 {len(self._template_map)} 个模板")
        
        # 直接在映射中查找
        exists = full_template_name in self._template_map
        logger.debug(f"模板 '{full_template_name}' 存在性检查结果: {exists}")
        
        if not exists and self._template_map:
            logger.debug(f"当前映射中的模板: {list(self._template_map.keys())}")
        
        return exists
        
    def _build_template_map(self):
        """
        内部方法：构建完整模板名称到SQL内容的映射
        用于支持快速查找功能
        """
        logger.debug("开始构建模板映射...")
        self._template_map.clear()
        
        # 遍历所有命名空间和模板，构建映射
        for namespace, templates in self.templates.items():
            logger.debug(f"处理命名空间: {namespace}，包含 {len(templates)} 个模板")
            for template_name, sql_content in templates.items():
                full_name = f"{namespace}.{template_name}"
                self._template_map[full_name] = sql_content
                logger.debug(f"添加映射: {full_name} -> {len(sql_content)} 字符的SQL")
        
        logger.debug(f"模板映射构建完成，总共 {len(self._template_map)} 个模板")
        logger.debug(f"映射中的模板列表: {list(self._template_map.keys())}")

    async def get_sql(self, sql_name):
        """
获取指定名称的SQL语句，支持命名空间格式：namespace.sql_name（异步版本）

        Args:
            sql_name: SQL模板名称，可以包含命名空间前缀

        Returns:
            str: SQL语句

        Raises:
            KeyError: 当找不到指定名称的SQL模板时
        """
        try:
            logger.debug(f"尝试获取SQL模板: {sql_name}")
            logger.debug(f"当前模板映射状态: {len(self._template_map) if self._template_map else 0} 个模板")
            
            # 快速检查模板是否存在
            if self._has_template(sql_name):
                logger.debug(f"在映射中找到模板: {sql_name}")
                # 如果在映射中存在，直接返回
                sql_content = self._template_map[sql_name]
                logger.debug(f"模板内容长度: {len(sql_content)} 字符")
                return sql_content
            
            # 如果内存中没有，检查Redis中是否有加载的模板数据
            templates_loaded = await self._check_templates_loaded_from_redis()
            if templates_loaded:
                await self._load_from_redis()
                # 再次尝试从内存中获取
                if self._has_template(sql_name):
                    sql_content = self._template_map[sql_name]
                    logger.debug(f"从Redis加载后找到模板: {sql_name}")
                    return sql_content
            
            logger.debug(f"映射中未找到模板: {sql_name}，尝试从原始模板中查找")
            
            # 尝试解析命名空间和SQL名称
            if '.' in sql_name:
                namespace, name = sql_name.split('.', 1)
            else:
                # 默认命名空间
                namespace = "default"
                name = sql_name

            logger.debug(f"解析结果 - 命名空间: {namespace}, 模板名: {name}")
            logger.debug(f"可用命名空间: {list(self.templates.keys())}")

            # 检查命名空间是否存在
            if namespace and namespace not in self.templates:
                logger.error(f"命名空间 '{namespace}' 不存在，可用命名空间: {list(self.templates.keys())}")
                raise KeyError(f"找不到命名空间: {namespace}")

            # 检查SQL名称是否存在
            available_templates = list(self.templates[namespace].keys())
            logger.debug(f"命名空间 '{namespace}' 下的可用模板: {available_templates}")

            if name and name not in self.templates[namespace]:
                logger.error(f"SQL模板 '{name}' 在命名空间 '{namespace}' 下不存在，可用模板: {available_templates}")
                raise KeyError(f"在命名空间 '{namespace}' 下找不到SQL模板: {name}")

            sql = self.templates[namespace][name]
            logger.debug(f"成功获取SQL模板，内容长度: {len(sql)} 字符")
            return sql

        except KeyError:
            raise
        except Exception as e:
            logger.error(f"获取SQL模板时出错: {str(e)}", exc_info=True)
            raise
            
    async def reload_template(self, template_path):
        """重新加载指定的SQL模板文件（异步版本）

        Args:
            template_path: SQL模板文件的路径，可以是绝对路径或相对于sql_dir的相对路径

        Returns:
            bool: 重新加载是否成功
        """
        try:
            # 确定模板文件的完整路径
            if self.sql_dir and not os.path.isabs(template_path):
                file_path = os.path.join(self.sql_dir, template_path)
            else:
                file_path = template_path

            # 获取文件的绝对路径
            file_abs_path = os.path.abspath(file_path)

            # 如果文件已加载，则从已加载列表中移除
            if file_abs_path in self._loaded_files:
                self._loaded_files.remove(file_abs_path)
                # 清除映射表，下次访问时会重新构建
                self._template_map.clear()
                # 清除Redis中的模板数据
                await self._clear_redis_data()

            # 重新加载模板，立即构建映射
            success = await self.load_template(file_abs_path, build_map=True)
            # 如果加载成功，更新Redis
            if success:
                await self._save_to_redis()
            return success
        except Exception as e:
            logger.error(f"重新加载SQL模板失败: {str(e)}", exc_info=True)
            return False
            
    async def _check_templates_loaded_from_redis(self) -> bool:
        """
        检查Redis中是否已经加载了SQL模板
        
        Returns:
            bool: 如果Redis中已加载返回True，否则返回False
        """
        try:
            return await redisKit.exists(REDIS_SQL_TEMPLATES_LOADED_KEY)
        except Exception as e:
            logger.error(f"检查Redis中SQL模板状态失败: {str(e)}")
            return False
            
    async def _load_from_redis(self):
        """
        从Redis加载SQL模板数据
        """
        try:
            # 加载templates
            templates_str = await redisKit.get_data(REDIS_SQL_TEMPLATES_KEY)
            if templates_str:
                self.templates = json.loads(templates_str)
                
            # 加载loaded_files
            loaded_files_str = await redisKit.get_data(REDIS_SQL_LOADED_FILES_KEY)
            if loaded_files_str:
                self._loaded_files = set(json.loads(loaded_files_str))
                
            # 加载template_map
            template_map_str = await redisKit.get_data(REDIS_SQL_TEMPLATE_MAP_KEY)
            if template_map_str:
                self._template_map = json.loads(template_map_str)
            
            # 无论从Redis加载了什么，都尝试重新构建一次映射以确保一致性
            if self.templates and not self._template_map:
                self._build_template_map()
                
        except Exception as e:
            logger.error(f"从Redis加载SQL模板数据失败: {str(e)}")
            # 出错时不抛出异常，使用内存中的默认值
            
    async def _save_to_redis(self):
        """
        将SQL模板数据保存到Redis
        """
        try:
            # 保存templates
            await redisKit.set_data(REDIS_SQL_TEMPLATES_KEY, json.dumps(self.templates))
            
            # 保存loaded_files
            await redisKit.set_data(REDIS_SQL_LOADED_FILES_KEY, json.dumps(list(self._loaded_files)))
            
            # 保存template_map
            await redisKit.set_data(REDIS_SQL_TEMPLATE_MAP_KEY, json.dumps(self._template_map))
            
            # 设置标记表示已加载
            await redisKit.set_data(REDIS_SQL_TEMPLATES_LOADED_KEY, "1")
            
        except Exception as e:
            logger.error(f"将SQL模板数据保存到Redis失败: {str(e)}")
            # 出错时不抛出异常，继续执行
            
    async def _clear_redis_data(self):
        """
        清除Redis中的模板数据
        """
        try:
            await redisKit.delete(REDIS_SQL_TEMPLATES_KEY)
            await redisKit.delete(REDIS_SQL_LOADED_FILES_KEY)
            await redisKit.delete(REDIS_SQL_TEMPLATE_MAP_KEY)
            await redisKit.delete(REDIS_SQL_TEMPLATES_LOADED_KEY)
        except Exception as e:
            logger.error(f"清除Redis中SQL模板数据失败: {str(e)}")
            # 出错时不抛出异常，继续执行