Fix custom registry assignment in model_config and inheritance

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

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"