paste-framework/paste/db/baseadapter.py

"""
数据访问适配器。主要封装了数据访问的基础方法。
"""
import datetime
import random
import uuid
from typing import Optional, Any, Union, List, Sequence

import pandas as pd
from sqlalchemy import Column, String, DateTime, Date, inspect, Table, Integer, select, Numeric, ForeignKey, func, \
    text, desc, Result, Engine, tuple_
from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, async_sessionmaker
from sqlalchemy.orm import sessionmaker, Session, InstanceState, InstrumentedAttribute, registry
from sqlalchemy.sql import Select
from sqlalchemy.sql.base import ColumnCollection
from sqlalchemy.sql.compiler import StrSQLCompiler
from sqlalchemy.sql.elements import BinaryExpression

from paste.db import engine
from paste.util import udict

LOCAL_DATE_FORMAT = '%Y-%m-%d'
LOCAL_TIME_FORMAT = '%H:%M:%S'
LOCAL_DATETIME_FORMAT = f'{LOCAL_DATE_FORMAT} {LOCAL_TIME_FORMAT}'


def guid() -> str:
    """
    生成 GUID。

    :return: GUID
    """
    return str(uuid.uuid5(uuid.NAMESPACE_DNS, str(uuid.uuid1()) + str(random.random()))).replace('-', '')


def get_session() -> Session:
    """
    获取线程安全的连接会话对象。

    :return: 连接会话对象
    """
    _engine: Engine = engine.connect_engine()
    session_maker = sessionmaker(bind=_engine)
    return session_maker()


def get_aio_session() -> AsyncSession:
    """
    获取异步连接会话对象。

    :return: 连接会话对象
    """
    _engine: AsyncEngine = engine.async_connect_engine()
    session_maker = async_sessionmaker(bind=_engine, class_=AsyncSession)
    return session_maker()


Base = registry().generate_base()
"""
取得 SQLAlchemy ORM 的声明式基类。
"""


class BaseAdapter(Base):
    """
    适配器类，集成数据库引擎、连接、会话、元数据操作、部分 DML 数据处理等功能。
    主要实现了同步执行方法。
    """
    __abstract__  = True

    def __init__(self, **kwargs):
        """
        构造模型对象。

        :param args:
        :param kwargs:
        """
        self._is_new = True
        super().__init__(**kwargs)

    @classmethod
    def get_session(cls):
        """
        线程安全的连接会话。

        :return: 同步会话对象
        """
        return get_session()

    @classmethod
    def get_aio_session(cls):
        """
        异步连接会话。

        :return: 异步会话对象
        """
        return get_aio_session()

    @classmethod
    def ping(cls):
        """
        尝试连接数据库服务器。

        :return: 连接结果
        """
        engine.connect_engine().connect()

    @classmethod
    async def tables_in_db(cls):
        """
        取得数据库中所有表名称。
        :return: 表名称列表
        """
        _query = select(text('table_name')).select_from(
            text('information_schema.tables')
        ).where(
            text('table_schema = database()')
        )
        _names = list()
        _session = cls.get_aio_session()
        try:
            _result: Result = await _session.execute(_query)
            _names = list(_result.scalars().all())
        finally:
            await _session.close()
        return _names

    @classmethod
    async def is_table_exist(cls, table_name: str):
        """
        判断表是否存在。
        :param table_name: 要判断的表名称
        """
        _query = select(func.count()).select_from(
            text('information_schema.tables')
        ).where(
            text('table_schema = database()'), text('table_name = :table_name')
        ).params(
            table_name=table_name
        )
        _count = 0
        _session = cls.get_aio_session()
        try:
            _result: Result = await _session.execute(_query)
            _count = int(_result.scalar())
        finally:
            await _session.close()
        return _count > 0

    @classmethod
    def table(cls) -> Table:
        """
        取得当前模型对象所对应的表对象。

        :return: 表对象
        """
        tables = cls.metadata.tables
        assert cls.__tablename__ is not None, '属性：cls.__tablename__ 未找到.'
        assert cls.__tablename__ in tables, ('数据库中未找到表：%s.' % cls.__tablename__)
        return tables[cls.__tablename__]

    @classmethod
    def columns(cls) -> List[Column]:
        """
        当前对象配置的所有列。

        :return: 所有列
        """
        cols = cls.table().columns
        assert isinstance(cols, ColumnCollection), '%s 列类型错误' % cls.__name__
        return list(cols)

    @classmethod
    def instrument_attr(cls, column_name: str) -> Optional[InstrumentedAttribute]:
        """
        取得查询条件绑定属性，用于配置查询条件表达式。

        :param column_name: 列名称
        :return: 类配置的列信息，即查询条件绑定属性
        """
        attr = getattr(cls, column_name)
        if not isinstance(attr, InstrumentedAttribute):
            return None
        return attr

    @classmethod
    def label(cls, column: Union[Column, Any]):
        """
        以字段的 comment 作为标签返回。

        :return: 指定列的备注
        """
        return column.comment

    @classmethod
    def labels(cls) -> dict:
        """
        取得所有列的名称描述。

        :return: 所有列的描述字典，结构为：{`field`: `comment`}
        """
        _label_dict = {}
        for _col in cls.columns():
            _label_dict[_col.key] = _col.comment
        return _label_dict

    @classmethod
    def field(cls, column: Union[Column, Any]) -> str:
        """
        取得字段名称。

        :param column: 列对象
        :return: 字段名称
        """
        return column.key

    @classmethod
    def fields(cls) -> List[str]:
        """
        取得所有字段名称列表。

        :return: 所有字段名称列表
        """
        return [col.key for col in cls.columns()]

    @classmethod
    def new_id(cls, **kwargs) -> str:
        """
        用GUID作为新的ID。

        :return: ID
        """
        return guid()

    @property
    def is_new(self):
        """
        通过检查是否已经存在ID判断是否为新建对象。

        :return: 是否为新建对象
        """
        return hasattr(self, '_is_new') and self._is_new

    @classmethod
    def raw_sql(cls, query) -> StrSQLCompiler:
        """
        显示编译后的原始 SQL 命令。

        :return: 编译后的 SQL 命令文本
        """
        raw_sql = query.compile(compile_kwargs={'literal_binds': True})
        return raw_sql

    @classmethod
    def row_count(cls, *where_expressions: BinaryExpression) -> int:
        """
        按条件取得记录行数。

        :param where_expressions: 查询条件
        :return: 返回记录行数
        """
        query = select(func.count(1)).select_from(cls).where(*where_expressions)
        session = cls.get_session()
        count = session.execute(query).scalars().first()
        session.close()
        return count

    @classmethod
    def exist_by_id(cls, id_val: Union[str, int]):
        """
        检查是否有存在的主键。

        :param id_val: 主键值
        """
        _wheres = [cls.instrument_attr('id') == id_val]
        count: int = cls.row_count(*_wheres)
        return count >= 1

    @classmethod
    def search_wheres(cls, likes: Optional[dict[str, str]] = None, **kwargs):
        """
        按参数组织查询条件。

        :param likes: 需要执行模糊查询的列
        :param kwargs: 属性填充参数
        :return: 查询条件列表
        """
        if likes is None:
            likes = []
        _query_model = cls().copy_from_dict(kwargs)
        return _query_model.filter_expressions(likes=likes)

    @classmethod
    def datalist_query(cls, row_list: List[dict], condition_cols: List[Column]) -> Optional[Select]:
        """
        根据数据列表、查询条件列生成查询。仅支持单表查询。

        :param row_list: 数据列表
        :param condition_cols: 要查询，且作为条件的列。注意次序与索引次序相同
        :return: 查询到的数据模型列表
        """
        if not row_list or not condition_cols:
            return None

        _condition_fields = [_col.key for _col in condition_cols]

        # 验证字段存在
        for field in _condition_fields:
            if not hasattr(cls, field):
                raise AttributeError(f"Field '{field}' not found in {cls.__name__}")

        # 构建值列表的元组形式
        _values_tuples = []
        for _row in row_list:
            _val_list = []
            for _f in _condition_fields:
                _val = udict.get_with_default(_row, _f, '')
                # 单独处理日期时间格式，但是此处固定格式
                if isinstance(_val, datetime.datetime):
                    _val = _val.strftime(LOCAL_DATETIME_FORMAT)
                if isinstance(_val, datetime.date):
                    _val = _val.strftime(LOCAL_DATE_FORMAT)
                _val_list.append(_val)

            # 过滤掉全为空的元组
            if any(v != '' for v in _val_list):
                _values_tuples.append(tuple(_val_list))

        if not _values_tuples:
            return None

        # 去重，提高性能
        _values_tuples = list(set(_values_tuples))

        # 使用参数化查询
        if len(condition_cols) == 1:
            # 单字段查询
            field = getattr(cls, _condition_fields[0])
            _single_values = [t[0] for t in _values_tuples]
            query = select(cls).where(field.in_(_single_values))
        else:
            # 多字段组合查询
            conditions = [getattr(cls, field) for field in _condition_fields]
            query = select(cls).where(tuple_(*conditions).in_(_values_tuples))

        return query

    @classmethod
    def dataframe_query(cls, row_df: pd.DataFrame, condition_cols: List[Column]):
        """
        根据数据列表、查询条件列生成查询。仅支持单表查询。

        :param row_df: 数据框架
        :param condition_cols: 要查询，且作为条件的列。注意次序与索引次序相同
        :return: 查询到的数据模型列表
        """
        if row_df.empty or not condition_cols:
            return None

        _condition_fields = [_col.key for _col in condition_cols]

        # 验证字段存在
        for field in _condition_fields:
            if not hasattr(cls, field):
                raise AttributeError(f"Field '{field}' not found in {cls.__name__}")

        # 去除重复行并删除包含 NaN 的行
        row_df = row_df[_condition_fields].drop_duplicates().dropna()

        if row_df.empty:
            return None

        if len(condition_cols) == 1:
            # 单字段查询
            _field_name = _condition_fields[0]
            field = getattr(cls, _field_name)
            # 直接获取 Series 的值
            _values_list = row_df[_field_name].tolist()
            if not _values_list:
                return None
            query = select(cls).where(field.in_(_values_list))
        else:
            # 多字段组合查询 - 使用向量化操作避免 iterrows()
            _values_tuples = [tuple(row) for row in row_df.values]
            # 再次去重确保
            _values_tuples = list(set(_values_tuples))

            if not _values_tuples:
                return None

            conditions = [getattr(cls, field) for field in _condition_fields]
            query = select(cls).where(tuple_(*conditions).in_(_values_tuples))

        return query

    @classmethod
    def find_by_id(cls, id_val: Union[str, int], reset_session: Optional[bool] = True) -> Optional['BaseAdapter']:
        """
        根据主键ID查找数据，确认有 id 字段后方可使用。

        :param id_val: 主键值
        :param reset_session: 重置会话连接
        :return: 数据模型对象
        """
        query = select(cls).where(cls.instrument_attr('id') == id_val)
        session = cls.get_session()
        model = session.execute(query).scalars().first()
        session.close()

        if reset_session and isinstance(model, BaseAdapter):
            model.reset_session()

        return model

    @classmethod
    def find_by_created(cls, from_t: datetime.datetime, to_t: Optional[datetime.datetime] = None) -> Sequence['BaseAdapter']:
        """
        按创建时间搜索。本方法使用了固定字段名称，确认有 created_at 字段后方可使用。

        :param from_t: 开始时间
        :param to_t: 结束时间
        """
        if to_t is None:
            to_t = from_t

        query = select(cls).where(cls.instrument_attr('created_at').between(from_t, to_t))
        session = cls.get_session()
        rows = session.execute(query).scalars().all()
        session.close()
        return rows

    @classmethod
    def find_by_updated(cls, from_t: datetime.datetime, to_t: Optional[datetime.datetime] = None) -> Sequence['BaseAdapter']:
        """
        按更新时间搜索。本方法使用了固定字段名称，确认有 updated_at 字段后方可使用。

        :param from_t: 开始时间
        :param to_t: 结束时间
        """
        if to_t is None:
            to_t = from_t

        query = select(cls).where(cls.instrument_attr('updated_at').between(from_t, to_t))
        session = cls.get_session()
        rows = session.execute(query).scalars().all()
        session.close()
        return rows

    @classmethod
    def find_by_datalist(cls, row_list: List[dict], condition_cols: List[Column]) -> Sequence['BaseAdapter']:
        """
        根据数据列表查询已经在数据库中的数据。

        :param row_list: 数据列表
        :param condition_cols: 要查询，且作为条件的列。注意次序与索引次序相同
        :return: 查询到的数据模型列表
        """
        rows: Sequence[cls] = []
        _query = cls.datalist_query(row_list=row_list, condition_cols=condition_cols)
        if _query is not None:
            session = cls.get_session()
            rows = session.execute(_query).scalars().all()
            session.close()

        return rows

    @classmethod
    def find_by_dataframe(cls, row_df: pd.DataFrame, condition_cols: List[Column]):
        """
        根据数据框架查询已经在数据库中的数据。

        :param row_df: 数据框架
        :param condition_cols: 要查询，且作为条件的列。注意次序与索引次序相同
        :return: 查询到的数据模型列表
        """
        _query = cls.dataframe_query(row_df=row_df, condition_cols=condition_cols)
        if _query is not None:
            session = cls.get_session()
            _rows = session.connection().execute(_query).all()
            session.close()
            return pd.DataFrame(_rows)

        return None

    def load(self, model_data: dict):
        """
        从字典对象载入数据。
        该方法仅从字典中读取与对象属性对应的数据，忽略其他数据。
        注意：与 copyXXX 方法不同，该方法跳过 model_data 中的 None 值，保留原始值不变

        :param model_data: 包含数据的字典对象
        """
        for attr in self.__dir__():
            if model_data.get(attr, None) is not None:
                self.__setattr__(attr, model_data.get(attr))
        return self

    def copy_from(self, source: 'BaseAdapter', mapping_fields: Optional[List[str]] = None):
        """
        从源数据模型复制数据，仅复制相同字段的数据。

        :param source: 源数据对象
        :param mapping_fields: 映射字段列表，若为空，则从源数据模型中获取字段列表
        """

        # 源对象字段列表
        if mapping_fields is None:
            mapping_fields = [column.key for column in source.columns()]

        for column in self.columns():
            # 若源对象中不包含同名字段，则跳过
            if column.key not in mapping_fields:
                continue

            # 取出对象属性值
            attr_val = source.ins_value(column.key)
            setattr(self, column.key, attr_val)

        return self

    def copy_from_dict(self, source: dict, mapping_keys: Union[list, tuple, set] = None, skip_none: bool = False):
        """
        从源字典对象复制数据，仅复制相同键值的数据。

        :param source: 源数据对象
        :param mapping_keys: 映射键列表，若为空，则从源字典对象中获取键列表
        :param skip_none: 是否跳过 None 值，即若源数据对象中属性值为 None 时，跳过
        :return 返回自身
        """

        # 源对象字段列表
        if mapping_keys is None:
            mapping_keys = source.keys()

        for column in self.columns():
            if column.key not in mapping_keys:
                # 若源对象中不包含同名字段，则跳过
                continue

            # 取出对象属性值
            attr_val = source.get(column.key, None)
            if skip_none and attr_val is None:
                # 若设置了跳过 None 且此时属性值为空时，跳过
                continue
            setattr(self, column.key, attr_val)

        return self

    def inspect(self) -> InstanceState:
        """
        取得对象的监视对象。

        :return: 监视对象
        """
        return inspect(self)

    def session(self) -> Session:
        """
        取得当前对象的连接会话对象。

        :return: 连接会话对象
        """
        return self.inspect().session

    def has_session(self) -> bool:
        """
        检测对象是否已有连接会话对象。

        :return: 有则返回 True 否则返回 False
        """
        return self.session() is not None

    def add_session(self):
        """
        将对象加入到连接会话对象。
        """
        if not self.has_session():
            self.get_session().add(self)
        return self.session()

    def close_session(self):
        """
        关闭会话。
        """
        if self.has_session():
            self.session().close()

    def reset_session(self):
        """
        重置对象的连接会话。
        """
        self.close_session()
        self.add_session()

    def ins_value(self, column_name: str) -> Any:
        """
        取得查询条件绑定的值，该值与对象属性一致。与::
            self.column_name
        或
            self['column_name']

        :param column_name: 列名
        :return: 对象属性值
        """
        # 取出对象属性值
        attr_val = getattr(self, column_name)
        if attr_val is None:
            return None
        return attr_val

    def filter_expressions(self, likes: Optional[dict[str, str]] = None) -> List[BinaryExpression]:
        """
        自动生成查询条件表达式。
        参数 likes 为需要执行模糊查询的字段字典::

            {
                field_name1: '%{}%',
                field_name2: '{}%',
                field_name3: '%{}',
                ...
            }

        :param likes: 需要执行模糊查询的字段字典
        :return: 包含查询条件表达式的列表
        """
        expressions = list()
        for column in self.columns():
            # 取出对象属性值
            attr_val = self.ins_value(column.key)
            # 若属性值为空，不增加条件，跳出
            if attr_val in (None, ''):
                continue

            # 取出类属性
            attr = self.instrument_attr(column.key)
            # 若类属性类型不正确，不增加条件，跳出
            if attr is None:
                continue

            # 取出列的数据类型
            column_type = column.type
            if isinstance(column_type, (String, Integer, Numeric, ForeignKey)):
                # 针对属性数据类型，进行不同处理
                if isinstance(attr_val, (list, tuple, set)):
                    if not attr_val:
                        # 跳过0长数组
                        continue
                    # 数组使用 in
                    expressions.append(attr.in_(attr_val))
                elif isinstance(attr_val, (int, float)):
                    # 数值类型
                    expressions.append(attr == attr_val)
                else:
                    # 字符类型，使用 like
                    if likes and attr.key in likes:
                        _f_str = likes[attr.key]
                        expressions.append(attr.like(_f_str.format(attr_val)))
                    else:
                        expressions.append(attr == attr_val)
            elif isinstance(column_type, (DateTime, Date)):
                if isinstance(attr_val, (list, tuple, set)):
                    if len(attr_val) == 2:
                        # 如果长度为 2 则使用 between and 方式
                        expressions.append(attr.between(attr_val[0], attr_val[1]))
                    else:
                        # 数组使用 in
                        expressions.append(attr.in_(attr_val))
                else:
                    expressions.append(attr.between(attr_val, attr_val))
            else:
                # 其他类型，用等号
                expressions.append(attr == attr_val)

        return expressions

    def gen_query(self, likes: Optional[dict[str, str]] = None):
        """
        根据自身参数生成查询对象。

        :param likes: 需要执行模糊查询的字段字典
        :return: 查询对象
        """
        cls = self.__class__
        expressions = self.filter_expressions(likes)
        _query: Select = select(cls).where(*expressions)
        return _query

    def find(self) -> Sequence['BaseAdapter']:
        """
        根据自身参数，查询数据库。

        :return: 查询到的结果对象
        """
        session = self.get_session()
        _model_list = session.execute(self.gen_query()).scalars().all()
        session.close()
        return _model_list

    def find_first(self) -> 'BaseAdapter':
        """
        根据自身参数，查询数据库，仅查询第一条。

        :return: 查询到的结果对象
        """
        session = self.get_session()
        _model = session.execute(self.gen_query()).scalars().first()
        session.close()
        return _model

    def find_piece(self, *where: BinaryExpression, offset=0, limit=500, is_desc=False,
                   likes: Optional[dict[str, str]] = None):
        """
        根据自身参数，查询数据库。

        :param where: 查询条件
        :param offset: 偏移量
        :param limit: 读取数量
        :param is_desc: 是否逆序排列
        :param likes: 模糊条件
        :return: 查询到的结果对象
        """
        clz = self.__class__
        expressions = self.filter_expressions(likes=likes)
        if where is not None:
            expressions += where

        query = select(clz).where(*expressions)
        if limit > 0:
            query = query.limit(limit=limit)
        if offset >= 0:
            query = query.offset(offset=offset)
        if is_desc:
            if hasattr(clz, 'id'):
                query = query.order_by(desc(clz.id))

        session = self.get_session()
        _model_list = session.execute(query).scalars().all()
        session.close()
        return _model_list

    def before_save(self):
        """
        保存前的动作，一般应当在子类中覆盖该方法，增加在保存前应当执行的动作。
        """
        if self.is_new:
            if hasattr(self, 'id'):
                setattr(self, 'id', self.new_id())

            if hasattr(self, 'created_at'):
                setattr(self, 'created_at', datetime.datetime.now().strftime(LOCAL_DATETIME_FORMAT))

        if hasattr(self, 'updated_at'):
            setattr(self, 'updated_at', datetime.datetime.now().strftime(LOCAL_DATETIME_FORMAT))

        return self

    def save(self, auto_expunge: Optional[bool] = True, session: Optional[Session] = None):
        """
        保存数据模型对象，若该对象尚未有连接会话，则自动加入连接会话。

        :param auto_expunge: 自动刷新对象并将其移出连接会话，不提供外部 session 时有效
        :param session: 会话对象，主要用于事务
        :return: 保存状态
        """
        _has_session: bool = True
        """
        该参数用于说明是否提供了外部 session 对象，默认为 True 时表示提供。
        """

        self.before_save()
        if session is None:
            _has_session = False
            session = self.add_session()

        try:
            session.add(self)
            if not _has_session:
                # 使用新会话时，直接提交
                session.commit()
            self._is_new = False
        except Exception as e:
            session.rollback()
            raise e
        else:
            if auto_expunge and not _has_session:
                session.refresh(self)
                session.expunge(self)
            return True

    def to_dict(self) -> dict:
        """
        数据模型转字典。递归处理内部类型对象。

        :return: 转换后的字典数据
        """
        # 模型数据字典
        m_dict = {}

        # 遍历处理内部转换
        for _key, _val in dict(self.__dict__).items():
            if f'{_key}'.startswith('_'):
                # 跳过所有私有属性
                continue

            if isinstance(_val, BaseAdapter):
                # 内部数据对象数据，直接转换
                m_dict[_key] = _val.to_dict()
            elif isinstance(_val, list):
                # 遍历转换数据对象列表
                _tmp_list = []
                for _i, _v in enumerate(_val):
                    if isinstance(_v, BaseAdapter):
                        _tmp_list.append(_v.to_dict())
                    else:
                        _tmp_list.append(_v)
                m_dict[_key] = _tmp_list
            elif isinstance(_val, dict):
                # 遍历转换数据对象字典
                _tmp_dict = {}
                for _ik, _iv in _val.items():
                    if isinstance(_iv, BaseAdapter):
                        _tmp_dict[_ik] = _iv.to_dict()
                    else:
                        _tmp_dict[_ik] = _iv
                m_dict[_key] = _tmp_dict
            else:
                # 其他属性直接赋值
                m_dict[_key] = _val

        return m_dict