from typing import ( Any, Dict, List, Set, Tuple, Type, cast, ) from sqlalchemy import ( inspect, ) from sqlalchemy.orm import ( Mapped, registry, relationship, ) from sqlalchemy.orm.decl_api import DeclarativeMeta from sqlmodel.main import ( BaseConfig, ModelField, ModelMetaclass, SQLModelConfig, Undefined, get_annotations, get_config_value, get_model_fields, get_relationship_to, is_table_model_class, set_config_value, __dataclass_transform__, Field, FieldInfo, RelationshipInfo, get_column_from_field, ) from typing_extensions import get_origin # Using CustomSQLModel and CustomSQLModelMetaclass, the default CustomSQLModelMetaclass will mark this sqlmodel # as __abstract__=True if you custom registry via kwargs, which will avoid it registering class into SQLAlchemy, # and then you will receive error while using some sqlalchemy feature. @__dataclass_transform__(kw_only_default=True, field_descriptors=(Field, FieldInfo)) class SQLModelMetaclass(ModelMetaclass, DeclarativeMeta): __sqlmodel_relationships__: Dict[str, RelationshipInfo] model_config: SQLModelConfig model_fields: Dict[str, FieldInfo] __config__: Type[SQLModelConfig] __fields__: Dict[str, ModelField] # type: ignore[assignment] # Replicate SQLAlchemy def __setattr__(cls, name: str, value: Any) -> None: if is_table_model_class(cls): DeclarativeMeta.__setattr__(cls, name, value) else: super().__setattr__(name, value) def __delattr__(cls, name: str) -> None: if is_table_model_class(cls): DeclarativeMeta.__delattr__(cls, name) else: super().__delattr__(name) # From Pydantic def __new__( cls, name: str, bases: Tuple[Type[Any], ...], class_dict: Dict[str, Any], **kwargs: Any, ) -> Any: relationships: Dict[str, RelationshipInfo] = {} dict_for_pydantic = {} original_annotations = get_annotations(class_dict) pydantic_annotations = {} relationship_annotations = {} for k, v in class_dict.items(): if isinstance(v, RelationshipInfo): relationships[k] = v else: dict_for_pydantic[k] = v for k, v in original_annotations.items(): if k in relationships: relationship_annotations[k] = v else: pydantic_annotations[k] = v dict_used = { **dict_for_pydantic, "__weakref__": None, "__sqlmodel_relationships__": relationships, "__annotations__": pydantic_annotations, } # Duplicate logic from Pydantic to filter config kwargs because if they are # passed directly including the registry Pydantic will pass them over to the # superclass causing an error allowed_config_kwargs: Set[str] = { key for key in dir(BaseConfig) if not ( key.startswith("__") and key.endswith("__") ) # skip dunder methods and attributes } config_kwargs = { key: kwargs[key] for key in kwargs.keys() & allowed_config_kwargs } new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs) new_cls.__annotations__ = { **relationship_annotations, **pydantic_annotations, **new_cls.__annotations__, } def get_config(name: str) -> Any: config_class_value = get_config_value( model=new_cls, parameter=name, default=Undefined ) if config_class_value is not Undefined: return config_class_value kwarg_value = kwargs.get(name, Undefined) if kwarg_value is not Undefined: return kwarg_value return Undefined config_table = get_config("table") if config_table is True: # If it was passed by kwargs, ensure it's also set in config set_config_value(model=new_cls, parameter="table", value=config_table) for k, v in get_model_fields(new_cls).items(): col = get_column_from_field(v) setattr(new_cls, k, col) # Set a config flag to tell FastAPI that this should be read with a field # in orm_mode instead of preemptively converting it to a dict. # This could be done by reading new_cls.model_config['table'] in FastAPI, but # that's very specific about SQLModel, so let's have another config that # other future tools based on Pydantic can use. set_config_value( model=new_cls, parameter="read_from_attributes", value=True ) # For compatibility with older versions # TODO: remove this in the future set_config_value(model=new_cls, parameter="read_with_orm_mode", value=True) config_registry = get_config("registry") if config_registry is not Undefined: config_registry = cast(registry, config_registry) # If it was passed by kwargs, ensure it's also set in config set_config_value(model=new_cls, parameter="registry", value=config_table) setattr(new_cls, "_sa_registry", config_registry) # noqa: B010 setattr(new_cls, "metadata", config_registry.metadata) # noqa: B010 setattr(new_cls, "__abstract__", False) # noqa: B010 <- Using True instead of False return new_cls # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models def __init__( cls, classname: str, bases: Tuple[type, ...], dict_: Dict[str, Any], **kw: Any ) -> None: # Only one of the base classes (or the current one) should be a table model # this allows FastAPI cloning a SQLModel for the response_model without # trying to create a new SQLAlchemy, for a new table, with the same name, that # triggers an error base_is_table = any(is_table_model_class(base) for base in bases) if is_table_model_class(cls) and not base_is_table: for rel_name, rel_info in cls.__sqlmodel_relationships__.items(): if rel_info.sa_relationship: # There's a SQLAlchemy relationship declared, that takes precedence # over anything else, use that and continue with the next attribute setattr(cls, rel_name, rel_info.sa_relationship) # Fix #315 continue raw_ann = cls.__annotations__[rel_name] origin = get_origin(raw_ann) if origin is Mapped: ann = raw_ann.__args__[0] else: ann = raw_ann # Plain forward references, for models not yet defined, are not # handled well by SQLAlchemy without Mapped, so, wrap the # annotations in Mapped here cls.__annotations__[rel_name] = Mapped[ann] # type: ignore[valid-type] relationship_to = get_relationship_to( name=rel_name, rel_info=rel_info, annotation=ann ) rel_kwargs: Dict[str, Any] = {} if rel_info.back_populates: rel_kwargs["back_populates"] = rel_info.back_populates if rel_info.link_model: ins = inspect(rel_info.link_model) local_table = getattr(ins, "local_table") # noqa: B009 if local_table is None: raise RuntimeError( "Couldn't find the secondary table for " f"model {rel_info.link_model}" ) rel_kwargs["secondary"] = local_table rel_args: List[Any] = [] if rel_info.sa_relationship_args: rel_args.extend(rel_info.sa_relationship_args) if rel_info.sa_relationship_kwargs: rel_kwargs.update(rel_info.sa_relationship_kwargs) rel_value = relationship(relationship_to, *rel_args, **rel_kwargs) setattr(cls, rel_name, rel_value) # Fix #315 # SQLAlchemy no longer uses dict_ # Ref: https://github.com/sqlalchemy/sqlalchemy/commit/428ea01f00a9cc7f85e435018565eb6da7af1b77 # Tag: 1.4.36 DeclarativeMeta.__init__(cls, classname, bases, dict_, **kw) else: ModelMetaclass.__init__(cls, classname, bases, dict_, **kw)