import json from typing import Literal, List, Any, Optional, Union, Generic, TypeVar, Dict, Type from pydantic import BaseModel from sqlalchemy import create_engine, Column, and_, asc, desc, func, String from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, Session, Query, Load from sqlalchemy_utils import database_exists, create_database, get_columns, get_column_key, get_type, get_primary_keys from pydantic import BaseModel from sqlalchemy import Column from sqlalchemy.orm import InstrumentedAttribute, ColumnProperty, Relationship, DeclarativeMeta, Session from sqlalchemy_utils import get_columns, get_column_key class SqlalchemyConnect: def __init__(self, Base: declarative_base, host="127.0.0.1", user="", password="", db="", db_type: Literal['mysql', 'postgresql'] = 'mysql'): self.Base = Base self.host = host self.user = user self.password = password self.db = db if db_type == 'mysql': self.engine = self.init_engine() if db_type == 'postgresql': self.engine = self.init_postgresql_engine() # self.async_engine = self.init_async_engine() def init_engine(self): engine = create_engine( f"mysql+pymysql://{self.user}:{self.password}@{self.host}/{self.db}?charset=utf8mb4", pool_recycle=3600 * 4) return engine def init_postgresql_engine(self): engine = create_engine(f'postgresql+psycopg2://{self.user}:{self.password}@{self.host}/{self.db}', pool_recycle=3600 * 4, echo=False) return engine def create_db(self): if not database_exists(self.engine.url): create_database(self.engine.url) def get_db(self) -> sessionmaker: try: session = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) db = session() yield db finally: db.close() def get_db_commit(self) -> sessionmaker: try: session = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) db = session() yield db db.commit() finally: db.close() def get_db_i(self) -> Session: # engine = create_engine( # f"mysql+pymysql://{self.user}:{self.password}@{self.host}/{self.db}?charset=utf8mb4") session = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) db = session() return db def init_database(self, create_db=False): if create_db: self.create_db() self.Base.metadata.create_all(bind=self.engine) def init_async_engine(self): async_engine = create_async_engine( f"mysql+aiomysql://{self.user}:{self.password}@{self.host}/{self.db}?charset=utf8mb4", pool_recycle=3600 * 4 ) return async_engine def get_model_column_keys(model): return [get_column_key(model, item) for item in get_columns(model)] def marge_col_names(col_names, include: List[str] = None, ex_include: List[str] = None): if include: col_names = [item for item in include if item in col_names] if ex_include: col_names = [item for item in col_names if item not in ex_include] return set(col_names) # 通用查询接口 QueryType = Literal['=', '==', '>', '>=', '<', '<=', 'in', 'like', 'range', "find_in_set"] class QueryParam(BaseModel): name: str type: QueryType value: Any class OrderParam(BaseModel): order_by: str type: Literal['asc', 'desc'] = 'asc' ModelInfo = TypeVar('ModelInfo', bound=BaseModel) class QueryInclude(BaseModel): include: Optional[List[str]] = None ex_include: Optional[List[str]] = None relation_use_id: bool = False # 为true的话，m2m关系返回的是id数组 class QueryParams(QueryInclude, Generic[ModelInfo]): params: Optional[List[QueryParam]] = None order: Optional[OrderParam] = None query: Optional[ModelInfo] = None page: int page_size: int def query_common(db: Session, model, query_params: QueryParams): """ 通用数据库查询接口 """ query = db.query(model) # if query_params.ex_include: # query = query.options(Load(model).load_only(*query_params.include)) # if query_params.include: # query = query.options(Load(model).noload(*query_params.ex_include)) page = query_params.page page_size = query_params.page_size # 排序 if query_params.order: query = query_order(model, query, query_params.order) if query_params.query: query = query_common_query_core(model, query, query_params.query) # 筛选 if query_params.params: query = query_common_params_core(model, query, query_params.params) count = query.count() return count, query, page, page_size def query_common_with_page(db: Session, model, query_params: QueryParams): count, query, page, page_size = query_common(db, model, query_params) query = query.offset((page - 1) * page_size).limit(page_size) return count, query def query_order(model, query: Query, order_param: OrderParam): """ 查询结果排序 """ if order_param: if order_param.order_by: column: Column = getattr(model, order_param.order_by) if order_param.type == 'asc': query = query.order_by(asc(column)) if order_param.type == 'desc': query = query.order_by(desc(column)) return query def query_common_params_core(model, query: Query, params: List[QueryParam]): cols: Dict[str, TypeInfo] = model.model_config.cols for item in params: name = item.name child_name = "" if '.' in name: [name, child_name] = name.split('.') query_type = item.type value = item.value col = cols[name] col_base_type = col.col_base_type column: Column = getattr(model, name) if col_base_type == 'relation' and col.relation: if col.relation.relation_type in ['m2o', 'o2o']: if col.relation.source_foreign_key: query = query.join(col.relation.target_model, getattr(col.relation.target_model, col.relation.target_id_key) == getattr(model, col.relation.source_foreign_key)) else: query = query.join(col.relation.target_model) column = getattr(col.relation.target_model, child_name) else: if col_base_type == 'json': column = column[child_name] if query_type in ['=', '==']: if col_base_type == 'relation' and col.relation and not child_name: relation_type = col.relation.relation_type secondary = col.relation.secondary target_secondary_key = col.relation.target_secondary_key source_secondary_key = col.relation.source_secondary_key if relation_type == 'm2m': query = query.join(secondary, model.id == secondary.columns.get(source_secondary_key)) \ .filter(secondary.columns.get(target_secondary_key) == value) else: if col.col_base_type == 'datetime': query = query.filter(func.date(column) == value) elif col_base_type == 'json': query = query.filter(column == json.dumps(value)) else: query = query.filter(column == value) elif query_type == ">": query = query.filter(column > value) elif query_type == "<": query = query.filter(column < value) elif query_type == ">=": query = query.filter(column >= value) elif query_type == "<=": query = query.filter(column <= value) elif query_type == "in": if col_base_type == 'relation' and col.relation: relation_type = col.relation.relation_type if relation_type == 'm2m': secondary = col.relation.secondary target_secondary_key = col.relation.target_secondary_key source_secondary_key = col.relation.source_secondary_key if value and type(value) == list: query = query.join(secondary, model.id == secondary.columns.get(source_secondary_key)) \ .filter(secondary.columns.get(target_secondary_key).in_(value)) else: query = query.join(secondary, model.id == secondary.columns.get(source_secondary_key)) \ .filter(secondary.columns.get(target_secondary_key) == value) else: query = query.filter(column.in_(value)) elif query_type == "like": query = query.filter(column.cast(String).like(f'%{value}%')) elif query_type == "find_in_set": query = query.filter(func.find_in_set(value, column)) elif query_type == "range": if value: if len(value) == 1 and value[0] is not None: query = query.filter(column >= value[0]) if len(value) == 2: if value[0] is not None and value[1] is None: query = query.filter(column >= value[0]) if value[0] is None and value[1] is not None: query = query.filter(column <= value[1]) if value[0] is not None and value[1] is not None: query = query.filter(and_(column >= value[0], column <= value[1])) return query def query_common_query_core(model, query: Query, query_data: BaseModel): q_data = query_data.dict(exclude_unset=True) cols: Dict[str, TypeInfo] = model.model_config.cols for key, value in q_data: annotation = query_data.__annotations__[key] if key in cols: column: Column = getattr(model, key) col = cols[key] col_base_type = col.col_base_type if col_base_type in ['int', 'float', 'any', 'enum', 'str']: if type(value) == list and value: query = query.filter(column.in_(value)) else: if col_base_type == 'str' and value: query = query.filter(column.like(f'%{value}%')) else: query = query.filter(column == value) elif col_base_type in ['date', 'datetime', 'time']: if value: if type(value) == list: if len(value) == 2: if value[0] is not None: query = query.filter(column >= value) if value[1] is not None: query = query.filter(column <= value) else: if col_base_type == 'datetime': # datetime转date来查 query = query.filter(func.date(column) == value) else: query = query.filter(column == value) elif col_base_type == 'relation' and col.relation: relation_type = col.relation.relation_type if relation_type == 'm2m': secondary = col.relation.secondary target_secondary_key = col.relation.target_secondary_key source_secondary_key = col.relation.source_secondary_key if value and type(value) == list: query = query.join(secondary, model.id == getattr(secondary, source_secondary_key)) \ .filter(getattr(secondary, target_secondary_key).in_(value)) else: query = query.join(secondary, model.id == getattr(secondary, source_secondary_key)) \ .filter(getattr(secondary, target_secondary_key) == value) elif relation_type == 'o2m': target_model = col.relation.target_model target_id_key = col.relation.target_id_key if value and type(value) == list: query = query.join(target_model).filter(getattr(target_model, target_id_key).in_(value)) else: query = query.join(target_model).filter(getattr(target_model, target_id_key == value)) return query def tree_table_to_json(db: Session, model, belong_str='belong', key_str='id'): """ 树结构数据库存储到json """ item_list = db.query(model).all() node_list = [] for item in item_list: item_dic = item.to_dict() item_dic['children'] = [] node_list.append(item_dic) nodes_dic = {node[key_str]: node for node in node_list} tree = [] for node in node_list: belong = node[belong_str] if belong: if belong in nodes_dic: nodes_dic[belong]['children'].append(node) else: tree.append(node) return tree def get_relation(col: InstrumentedAttribute): prop = col.property if isinstance(prop, Relationship): model = prop.entity.entity if prop.secondary is not None: columns = prop.secondary.columns foreign_key_0 = list(columns[0].foreign_keys)[0] foreign_key_1 = list(columns[1].foreign_keys)[0] if foreign_key_0.column.table.key == model.__tablename__: target_secondary_key = columns[0].key target_key = foreign_key_0.column.key source_secondary_key = columns[1].key source_key = foreign_key_1.column.key else: source_secondary_key = columns[0].key source_key = foreign_key_0.column.key target_secondary_key = columns[1].key target_key = foreign_key_1.column.key return {'relation': 'm2m', 'model': model, 'secondary': prop.secondary, 'source_secondary_key': source_secondary_key, 'source_key': source_key, 'target_secondary_key': target_secondary_key, 'target_key': target_key, } else: return {'relation': 'm2o', 'model': model} else: return False def get_relation_col_dic(model: Type[DeclarativeMeta]): relation_col_dic = { } for col_name in dir(model): col = getattr(model, col_name) if isinstance(col, InstrumentedAttribute): relation = get_relation(col) if relation: relation_col_dic[col_name] = relation return relation_col_dic # 一个模型的字段类型 # 基本类型 RelationType = Literal['m2m', 'm2o', 'o2m', 'o2o'] BaseType = Literal[ 'int', 'float', 'bool', 'str', 'date', 'datetime', 'list', 'time', 'any', 'enum', 'set', 'relation', "json", "jsonb"] class Relation: relation_type: RelationType = None source_model: Type[DeclarativeMeta] = None # 左表 target_model: Type[DeclarativeMeta] = None # 右表 secondary: Any = None # 中间表 source_secondary_key: str = None # 中间表左表字段 source_key: str = None # 左表在中间表使用的key target_secondary_key: str = None # 中间表右表字段 target_key: str = None # 右表在中间表使用的key source_id_key: str = None # 左表id target_id_key: str = None # 右表id source_foreign_key: str = None # 左表里字段对应的foreign_key字段 ColOrgTypeEnum = Literal[ "BigInteger", "Boolean", "Date", "DateTime", "Enum", "Double", "Float", "Integer", "Interval", "LargeBinary", "MatchType", "Numeric", "PickleType", "SchemaType", "SmallInteger", "String", "Text", "Time", "Unicode", "UnicodeText", "Uuid", "JSON", "JSONB", 'DeclarativeMeta', ] class TypeInfo: col_org_type: ColOrgTypeEnum = None col_base_type: BaseType = None relation: Relation = None class ModelConfig: cols: Dict[str, TypeInfo] = {} id_key: str = 'id' COL_ORG_TYPE_MAP = { "BigInteger": 'int', "Boolean": 'bool', "Date": 'date', "DateTime": 'datetime', "Enum": 'enum', "Double": 'float', "Float": 'float', "Integer": 'int', "Interval": 'any', "LargeBinary": 'any', "MatchType": 'any', "Numeric": 'any', "PickleType": 'any', "SchemaType": 'any', "SmallInteger": 'int', "String": 'str', "Text": 'str', "Time": 'time', "Unicode": 'str', "UnicodeText": 'str', "Uuid": 'str', "JSON": "json", "JSONB": "jsonb", 'DeclarativeMeta': 'relation' } def get_model_id_key(model: Type[DeclarativeMeta]): id_keys = list(get_primary_keys(model).keys()) print(id_keys) if id_keys: return id_keys[0] def get_relation_info(col: InstrumentedAttribute) -> Relation: prop = col.property if isinstance(prop, Relationship): r = Relation() source_model = col.property.parent.entity source_id_key = get_model_id_key(source_model) target_model = prop.entity.entity target_id_key = get_model_id_key(target_model) source_uselist = prop.uselist target_uselist = False r.source_model = source_model r.target_model = target_model r.source_id_key = source_id_key r.target_id_key = target_id_key if col.property._user_defined_foreign_keys: r.source_foreign_key = list(col.property._user_defined_foreign_keys)[0].name print(r.source_foreign_key) # 找到另一个表内反向引用的字段id for name in get_model_col_names(target_model): target_col = getattr(target_model, name) col_org_type = get_type(target_col) if source_model == col_org_type: target_uselist = target_col.property.uselist # 寻找是否右表有对应的foreign_key if hasattr(target_col, 'foreign_keys'): for foreign_key in list(target_col.foreign_keys): if foreign_key.column.table.name == source_model.__tablename__: r.source_key = foreign_key.column.key # 找到本表内反向引用的字段id # for name in get_model_col_names(source_model): # source_col = getattr(target_model, name) # col_org_type = get_type(source_col) # if target_model == col_org_type: # source_uselist = source_col.property.uselist # # 寻找是否右表有对应的foreign_key # if hasattr(source_col,'foreign_keys'): # for foreign_key in List[source_col.foreign_keys]: # if foreign_key.column.table.name==target_model.__tablename__: # r.target_key=foreign_key.column.key if (source_uselist and target_uselist) or prop.secondary is not None: r.relation_type = 'm2m' columns = prop.secondary.columns foreign_key_0 = list(columns[0].foreign_keys)[0] foreign_key_1 = list(columns[1].foreign_keys)[0] if foreign_key_0.column.table.key == target_model.__tablename__: target_secondary_key = columns[0].key target_key = foreign_key_0.column.key source_secondary_key = columns[1].key source_key = foreign_key_1.column.key else: source_secondary_key = columns[0].key source_key = foreign_key_0.column.key target_secondary_key = columns[1].key target_key = foreign_key_1.column.key r.secondary = prop.secondary r.source_secondary_key = source_secondary_key r.target_secondary_key = target_secondary_key r.source_key = source_key r.target_key = target_key elif source_uselist == True and target_uselist == False: r.relation_type = 'o2m' elif source_uselist == False and target_uselist == False: r.relation_type = 'o2o' elif source_uselist == False and target_uselist == True: r.relation_type = 'm2o' return r else: return None def get_model_config(model: Type[DeclarativeMeta]) -> ModelConfig: """ 这个接口获取一个sql模型的各种信息 :param model: :return: """ # 映射 model_config = ModelConfig() cols = {} for col_name in dir(model): col = getattr(model, col_name) if isinstance(col, InstrumentedAttribute): if hasattr(col, 'primary_key') and col.primary_key: model_config.id_key = col_name type_info = TypeInfo() col_org_type = get_type(col).__class__.__name__ col_base_type = COL_ORG_TYPE_MAP.get(col_org_type) or 'any' relation = get_relation_info(col) type_info.col_org_type = col_org_type type_info.col_base_type = col_base_type type_info.relation = relation cols[col_name] = type_info model_config.cols = cols return model_config def get_model_col_names(model: Type[DeclarativeMeta]): columns = [] for col_name in dir(model): if isinstance(getattr(model, col_name), InstrumentedAttribute): columns.append(col_name) return columns class SalBase: __model_config: ModelConfig = None @classmethod @property def model_config(cls) -> ModelConfig: if cls.__model_config is None: cls.__model_config = 'process_ing' cls.__model_config = get_model_config(cls) return cls.__model_config def to_dict(self, include: List[str] = None, ex_include: List[str] = None): col_names = marge_col_names([c.name for c in self.__table__.columns], include, ex_include) data = {name: getattr(self, name) for name in col_names} return data def to_full_dict(self, include: List[str] = None, ex_include: List[str] = None, relation_use_id=False): col_names = marge_col_names(set(self.model_config.cols.keys()), include, ex_include) data = {} for col_name in col_names: col = self.model_config.cols[col_name] col_val = getattr(self, col_name) relation = col.relation if relation: if relation.relation_type == 'm2m': if relation_use_id: col_data = [getattr(item, relation.target_id_key) for item in col_val] else: col_data = [item.to_dict() for item in col_val] data[col_name] = col_data elif relation.relation_type == 'o2m': if relation_use_id: col_data = [getattr(item, relation.target_id_key) for item in col_val] else: col_data = [item.to_dict() for item in col_val] data[col_name] = col_data elif relation.relation_type == 'o2o': data[col_name] = col_val and col_val.to_dict() elif relation.relation_type == 'm2o': data[col_name] = col_val and col_val.to_dict() else: data[col_name] = col_val else: data[col_name] = col_val return data class TableModel(DeclarativeMeta, SalBase): pass class ModelCRUD: def __init__(self, model: Type[TableModel]): self.model = model def update(self, db: Session, data: BaseModel): model = self.model id_key = self.model.model_config.id_key item = db.query(model).filter(getattr(model, id_key) == (getattr(data, id_key))).first() data = data.dict(exclude_unset=True) col_names = get_model_col_names(model) model_config: ModelConfig = model.model_config for k, v in data.items(): if '.' in k: paths = k.split('.') if paths[0] in col_names: col_type_info: TypeInfo = model.model_config.cols[paths[0]] if col_type_info.relation: # 一对一时，{"config.age":1} relation = col_type_info.relation if relation.relation_type in ["o2o"]: col = getattr(item, k) if not col: if v is not None: setattr(item, paths[0], relation.target_model(**{paths[1]: v})) else: setattr(col, paths[1], v) if k in col_names: col_type_info: TypeInfo = model.model_config.cols[k] # col = getattr(model, k) relation = col_type_info.relation if relation: if relation.relation_type in ['m2m']: if v is not None: if v and type(v[0]) == dict: v = [item[model_config.id_key] for item in v] # 多对多更新的值是一个id数组，所以先用id查出对应项，再更新 target_model = relation.target_model target_key = relation.target_id_key update_values = {child for child in db.query(target_model).filter(getattr(target_model, target_key).in_(v))} setattr(item, k, update_values) elif relation.relation_type == 'o2o': if type(v) == dict: col = getattr(item, k) if not col: setattr(item, k, relation.target_model(**v)) else: for v_k, v_v in v.items(): setattr(col, v_k, v_v) else: setattr(item, k, v) db.commit() db.refresh(item) return item def get(self, db: Session, data: BaseModel): item = db.query(self.model).filter(getattr(self.model, self.model.model_config.id_key) == ( getattr(data, self.model.model_config.id_key))).first() return item def add(self, db: Session, data: BaseModel): model = self.model item = model() data = data.dict(exclude_unset=True) col_names = get_model_col_names(model) for k, v in data.items(): if k in col_names: col_type_info: TypeInfo = model.model_config.cols[k] # col = getattr(model, k) relation = col_type_info.relation if relation: if relation.relation_type == 'm2m': if v is not None: target_model = relation.target_model update_values = {child for child in db.query(target_model).filter( getattr(target_model, relation.target_id_key).in_(v))} setattr(item, k, update_values) elif relation.relation_type == 'o2o': if type(v) == dict: setattr(item, k, relation.target_model(**v)) else: setattr(item, k, v) db.add(item) db.commit() db.refresh(item) return item def query(self, db: Session, data: QueryParams) -> (int, List[Any]): return query_common(db, self.model, data) def delete(self, db: Session, data: BaseModel): db.query(self.model).filter(getattr(self.model, self.model.model_config.id_key) == ( getattr(data, self.model.model_config.id_key))).delete() db.commit() return True