VWED_server/services/online_script/built_in_modules/data_module.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
VWED.data 模块 - 数据存储和缓存
"""

import json
import uuid
from typing import Dict, Any, Optional, List
from sqlalchemy import text
from utils.logger import get_logger


class VWEDDataModule:
    """VWED.data 模块 - 数据存储和缓存"""
    
    def __init__(self, script_id: str):
        self.script_id = script_id
        self._storage: Dict[str, Any] = {}
        self.logger = get_logger(f"services.online_script.data_module.{script_id}")
    
    def get(self, key: str, default=None):
        """获取数据"""
        return self._storage.get(key, default)
    
    def set(self, key: str, value: Any):
        """设置数据"""
        self._storage[key] = value
    
    def delete(self, key: str):
        """删除数据"""
        if key in self._storage:
            del self._storage[key]
    
    def clear(self):
        """清空所有数据"""
        self._storage.clear()
    
    def keys(self):
        """获取所有键"""
        return list(self._storage.keys())
    
    def has(self, key: str) -> bool:
        """检查键是否存在"""
        return key in self._storage

    async def get_cache_param(self, key: str) -> Optional[str]:
        """
        获取缓存数据
        
        Args:
            key: 缓存键
            
        Returns:
            str | None: 字符串类型，JSON格式。如果数据不存在，将返回None
            
        Raises:
            Exception: 本方法会抛出异常
        """
        from sqlalchemy import select
        from data.models.datacachesplit import VWEDDataCacheSplit
        from data.session import get_async_session
        
        async with get_async_session() as session:
            query = select(VWEDDataCacheSplit).where(
                VWEDDataCacheSplit.data_key == key,
                VWEDDataCacheSplit.is_deleted == 0
            )
            result = await session.execute(query)
            record = result.scalar_one_or_none()
            
            if not record or not record.data_value:
                return None
            
            return record.data_value

    async def put_cache_param(self, key: str, value: str) -> None:
        """
        缓存数据
        
        将参数中的key和value键值对在系统中进行缓存，缓存的数据将持久化，
        之后可通过get_cache_param方法根据key获取对应的value。
        
        Args:
            key: 缓存键
            value: 缓存值。如果要缓存一个对象，需要先序列化为JSON字符串
            
        Returns:
            None
            
        Note:
            本方法不会抛出异常
        """
        try:
            from sqlalchemy import select, insert, update
            from data.models.datacachesplit import VWEDDataCacheSplit
            from data.session import get_async_session
            
            async with get_async_session() as session:
                # 查询是否已存在相同key的记录且未被删除
                query = select(VWEDDataCacheSplit).where(
                    VWEDDataCacheSplit.data_key == key,
                    VWEDDataCacheSplit.is_deleted == 0
                )
                result = await session.execute(query)
                existing_active_record = result.scalar_one_or_none()
                
                # 如果有未删除的记录，则更新它
                if existing_active_record:
                    stmt = update(VWEDDataCacheSplit).where(
                        VWEDDataCacheSplit.id == existing_active_record.id
                    ).values(
                        data_value=json.dumps(value) if not isinstance(value, str) else value,
                    )
                    await session.execute(stmt)
                else:
                    # 如果没有未删除的记录，创建新记录
                    stmt = insert(VWEDDataCacheSplit).values(
                        id=str(uuid.uuid4()),
                        data_key=key,
                        data_value=json.dumps(value) if not isinstance(value, str) else value,
                        is_deleted=0
                    )
                    await session.execute(stmt)
                
                await session.commit()
        except Exception:
            # 本方法不会抛出异常，静默处理
            pass

    async def clear_cache_param(self, key: str) -> None:
        """
        删除缓存数据
        
        Args:
            key: 缓存键
            
        Returns:
            None
            
        Note:
            本方法不抛出异常
        """
        try:
            from sqlalchemy import update, select
            from data.models.datacachesplit import VWEDDataCacheSplit
            from data.session import get_async_session
                    
            async with get_async_session() as session:
                # 首先检查是否存在该条记录
                query = select(VWEDDataCacheSplit).where(
                    VWEDDataCacheSplit.data_key == key,
                    VWEDDataCacheSplit.is_deleted == 0
                )
                result = await session.execute(query)
                exists = result.scalar_one_or_none() is not None
                
                if exists:
                    # 软删除具有指定键的记录
                    stmt = update(VWEDDataCacheSplit).where(
                        VWEDDataCacheSplit.data_key == key,
                        VWEDDataCacheSplit.is_deleted == 0
                    ).values(
                        is_deleted=1,
                    )
                    await session.execute(stmt)
                    await session.commit()
        except Exception:
            # 本方法不抛出异常，静默处理
            pass

    async def get_all_cache_params(self) -> str:
        """
        获取所有的缓存数据
        
        获取缓存块的全部缓存数据，并以JSON格式返回。
        
        Returns:
            str: 字符串类型，JSON格式。如果数据不存在，将返回空字符串
            
        Raises:
            Exception: 本方法会抛出异常
        """
        from sqlalchemy import select
        from data.models.datacachesplit import VWEDDataCacheSplit
        from data.session import get_async_session
        
        async with get_async_session() as session:
            query = select(VWEDDataCacheSplit).where(
                VWEDDataCacheSplit.is_deleted == 0
            )
            result = await session.execute(query)
            records = result.scalars().all()
            
            if not records:
                return ""
            
            # 构建缓存数据字典
            cache_data = {}
            for record in records:
                try:
                    # 尝试将值解析为JSON，如果失败则使用原始字符串
                    cache_data[record.data_key] = json.loads(record.data_value)
                except (json.JSONDecodeError, TypeError):
                    cache_data[record.data_key] = record.data_value
            
            return json.dumps(cache_data, ensure_ascii=False)

    def get_settings_config_value(self, key: str) -> str:
        """获取 settings.py 的配置
        
        Args:
            key: 配置字段名
            
        Returns:
            str: JSON 字符串格式的配置值
        """
        try:
            from config.settings import settings
            
            # 使用 getattr 获取配置值
            config_value = getattr(settings, key, None)
            if config_value is not None:
                return json.dumps(config_value, ensure_ascii=False)
            
            self.logger.warning(f"[{self.script_id}] 未找到配置项: {key}")
            return "{}"
            
        except Exception as e:
            self.logger.error(f"[{self.script_id}] 获取设置配置失败: {e}")
            return "{}"

    async def jdbc_execute_sql(self, sql: str) -> bool:
        """执行 SQL 语句
        
        通过 SQL 命令操作 RDS 系统数据库，可用于创建自定义数据表。
        
        Args:
            sql: SQL 语句
            
        Returns:
            bool: True表示执行成功，False表示执行失败
        """
        try:
            from data.session import get_async_session
            
            async with get_async_session() as session:
                await session.execute(text(sql))
                await session.commit()
                
            self.logger.info(f"[{self.script_id}] SQL执行成功: {sql[:100]}...")
            return True
            
        except Exception as e:
            self.logger.error(f"[{self.script_id}] SQL执行失败: {e}, SQL: {sql}")
            return False

    async def jdbc_query(self, sql: str) -> str:
        """使用 SQL 查询数据
        
        Args:
            sql: SQL 查询语句
            
        Returns:
            str: JSON格式的查询结果
        """
        try:
            from data.session import get_async_session
            
            async with get_async_session() as session:
                result = await session.execute(text(sql))
                
                # 获取列名
                columns = list(result.keys()) if result.keys() else []
                
                # 获取所有行数据
                rows = result.fetchall()
                
                # 转换为字典列表
                data = []
                for row in rows:
                    row_dict = {}
                    for i, col in enumerate(columns):
                        row_dict[col] = str(row[i]) if row[i] is not None else None
                    data.append(row_dict)
                
                return json.dumps(data, ensure_ascii=False)
                
        except Exception as e:
            self.logger.error(f"[{self.script_id}] SQL查询失败: {e}, SQL: {sql}")
            return "[]"

    async def jdbc_insert_or_update(self, sql: str, *sql_params: Any) -> int:
        """更新数据
        
        更新数据记录，支持动态参数。
        
        Args:
            sql: SQL语句，支持参数占位符 ?
            sql_params: 动态参数列表
            
        Returns:
            int: 0表示更新失败，大于0表示更新的记录数目
        """
        try:
            from data.session import get_async_session
            
            # 将 ? 占位符转换为 :param_n 格式
            param_sql = sql
            params = {}
            for i, param in enumerate(sql_params):
                param_key = f"param_{i}"
                param_sql = param_sql.replace("?", f":{param_key}", 1)
                params[param_key] = param
            
            async with get_async_session() as session:
                result = await session.execute(text(param_sql), params)
                await session.commit()
                
                affected_rows = result.rowcount
                self.logger.info(f"[{self.script_id}] SQL更新成功，影响行数: {affected_rows}")
                return affected_rows
                
        except Exception as e:
            self.logger.error(f"[{self.script_id}] SQL更新失败: {e}, SQL: {sql}, 参数: {sql_params}")
            return 0

    async def jdbc_query_count(self, sql: str) -> Optional[int]:
        """查询数据数量
        
        执行自定义 SQL 语句查询满足条件数据的数量。
        注意：只能使用 select count 语句。
        
        Args:
            sql: SQL count查询语句
            
        Returns:
            Optional[int]: None表示查询失败，其他数字表示查询到的记录数目
        """
        try:
            # 简单验证SQL是否为count查询
            sql_lower = sql.lower().strip()
            if not sql_lower.startswith("select count"):
                self.logger.error(f"[{self.script_id}] 只允许使用 select count 语句: {sql}")
                return None
            
            from data.session import get_async_session
            
            async with get_async_session() as session:
                result = await session.execute(text(sql))
                count = result.scalar()
                
                return int(count) if count is not None else 0
                
        except Exception as e:
            self.logger.error(f"[{self.script_id}] Count查询失败: {e}, SQL: {sql}")
            return None