diff --git a/sqlmodel/main.py b/sqlmodel/main.py index 40fe64e423..9064baa937 100644 --- a/sqlmodel/main.py +++ b/sqlmodel/main.py @@ -654,14 +654,20 @@ def get_config(name: str) -> Any: # TODO: remove this in the future new_cls.model_config["read_with_orm_mode"] = True # ty: ignore[invalid-key] - 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 - new_cls.model_config["registry"] = config_table - setattr(new_cls, "_sa_registry", config_registry) # noqa: B010 - setattr(new_cls, "metadata", config_registry.metadata) # noqa: B010 - setattr(new_cls, "__abstract__", True) # noqa: B010 + has_custom_registry = "registry" in kwargs or ( + "model_config" in class_dict and "registry" in class_dict["model_config"] + ) + if has_custom_registry: + config_registry = kwargs.get("registry", Undefined) + if config_registry is Undefined: + config_registry = class_dict["model_config"].get("registry", Undefined) + 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 + new_cls.model_config["registry"] = config_registry + setattr(new_cls, "_sa_registry", config_registry) # noqa: B010 + setattr(new_cls, "metadata", config_registry.metadata) # noqa: B010 + setattr(new_cls, "__abstract__", True) # noqa: B010 return new_cls # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models diff --git a/tests/test_main.py b/tests/test_main.py index fa40b71853..96c15d1dbe 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -216,3 +216,37 @@ class Hero(SQLModel, table=True): assert len(foreign_keys) == 1 assert foreign_keys[0].ondelete == "CASCADE" assert team_id_column.nullable is False + + +def test_custom_registry_in_model_config(clear_sqlmodel): + from sqlalchemy.orm import registry + + custom_reg = registry() + + class CustomBase(SQLModel, registry=custom_reg): + pass + + assert CustomBase.model_config["registry"] is custom_reg + assert getattr(CustomBase, "_sa_registry") is custom_reg # noqa: B009 + assert CustomBase.metadata is custom_reg.metadata + + class Hero(CustomBase, table=True): + id: int | None = Field(default=None, primary_key=True) + name: str + + assert Hero.model_config["registry"] is custom_reg + assert getattr(Hero, "_sa_registry") is custom_reg # noqa: B009 + assert Hero.metadata is custom_reg.metadata + assert hasattr(Hero, "__mapper__") + assert Hero.__table__.name == "hero" + + engine = create_engine("sqlite://") + custom_reg.metadata.create_all(engine) + + with Session(engine) as session: + hero = Hero(name="Deadpond") + session.add(hero) + session.commit() + session.refresh(hero) + assert hero.id is not None + assert hero.name == "Deadpond"