利用pydantic model封装db table的多过滤条件

我们在筛选 db table 数据的时候，经常需要支持多个字段作为过滤条件

在golang里面，我们通常定义一个 struct XxxFilter，通过一个结构体的非默认字段来指定查询条件，这很直观

在python里面，我们可以利用 pydantic model 模仿这样的实现

from pydantic import BaseModel
from sqlalchemy.orm.query import Query
from app.extension.plugin import db
# db is an instance of SQLAlchemy


class FilterOperator:
    """
    定义 DynamicFilter 支持的过滤操作符 （常量）
    """
    LIKE = "op_like"
    IN = "op_in"
    EQ = "op_eq"

    @classmethod
    def supported_operators(cls) -> list[str]:
        """
        返回所有支持的过滤操作符
        """
        return [cls.LIKE, cls.IN, cls.EQ]


class DynamicFilter(BaseModel):
    """
    This class is used to create a dynamic filter 
    作为具体过滤条件的结构体应该继承该类

    DynamicFilter 支持 [等值查询，模糊查询，in查询] 过滤

    等值查询 需将 DynamicFilter 字段的 description 设置为 FilterOperator.EQ
    模糊查询 需将 DynamicFilter 字段的 description 设置为 FilterOperator.LIKE
    in查询 需将 DynamicFilter 字段的 description 设置为 FilterOperator.IN

    如果 DynamicFilter 字段无 description ，则：
    - list/set/tuple 对象默认为 FilterOperator.IN
    - 其他对象默认为 FilterOperator.EQ

    usage example:
    ```
    class XxxFilter(DynamicFilter):
        id: int = Field(None, description=FilterOperator.EQ)
        name: list[str] = Field(None, description=FilterOperator.IN)
    ````
    表示: 待查表支持 id的等值过滤 和 name的in过滤
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # validate每个field的description
        for field_name, field in self.__fields__.items():

            field_des = field.field_info.description

            if not field_des or field_des not in FilterOperator.supported_operators():
                field_value = self.dict()[field_name]
                if isinstance(field_value, list) or \
                        isinstance(field_value, set) or \
                        isinstance(field_value, tuple):
                    # list/set/tuple 对象默认使用 FilterOperator.IN 操作符
                    field.field_info.description = FilterOperator.IN
                else:
                    # 其他对象默认使用 FilterOperator.EQ 操作符
                    field.field_info.description = FilterOperator.EQ


class DynamicFilterBuilder:

    @staticmethod
    def add_dynamic_filter_to_query(base_query: Query, table_to_filter: db.Model, filter: DynamicFilter) -> Query:
        """
        @desc    :  根据过滤条件filter和table，排除filter中的None值，自动构建新的orm query with filter 
        ---------------
        :param base_query: 基础Query
        :param table_to_filter: 指定过滤条件作用的目标table
        :param filter: 指定过滤条件
        :return : Query with filter
        """

        for field_name, field in filter.__fields__.items():
            table_attr = getattr(table_to_filter, field_name, None)

            if table_attr is not None:
                filter_value = getattr(filter, field_name, None)
                if filter_value is None:
                    continue

                filter_op = field.field_info.description

                if filter_op == FilterOperator.EQ:
                    base_query = base_query.filter(table_attr == filter_value)
                elif filter_op == FilterOperator.IN:
                    base_query = base_query.filter(table_attr.in_(filter_value))
                elif filter_op == FilterOperator.LIKE:
                    base_query = base_query.filter(table_attr.like(f"%{filter_value}%"))
                else:
                    raise Exception(f"Unsupported search filter operator: {filter_op}")

        return base_query

以下是简单的使用示例:

# example usage of DynamicFilter
from pydantic import Field
from sqlalchemy import create_engine, Column, Integer, String, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from .dynamic_filter import FilterOperator, DynamicFilter, DynamicFilterBuilder

Base = declarative_base()

# 定义表结构
class Employee(Base):
    __tablename__ = 'employees'
    id = Column(Integer, primary_key=True)
    name = Column(String)


# 定义表的查询条件
class EmployeeFilter(DynamicFilter):
    """
    约定：
    - 每个 Field 默认值为 None，以方便生成sql时剔除
    - 在description指定过滤操作符
    """
    id: int = Field(None, description=FilterOperator.EQ)
    name: list[str] = Field(None, description=FilterOperator.IN)



if __name__ == '__main__':
    engine = create_engine('mysql+pymysql://user:password@dbhost/mydatabase')
    Session = sessionmaker(bind=engine) # Session is a sessionmaker instance
    session = Session()
    
  # generate a base query
    query = session.query(Employee)
    
  # add filter to query
    filter_for_search = EmployeeFilter(name=["Bob", "Jimmy"])
    query = DynamicFilterBuilder.add_dynamic_filter_to_query(query, Employee, filter_for_search)
    print(query.all())

    

这样一来，我们就借助DynamicFilterBuilder.add_dynamic_filter_to_query，很方便地给任何Query对象增加过滤条件