Add support for implicit models when using a single storage

Author: frankie567

.zed/settings.json

@@ -18,14 +18,5 @@
                 }
             ]
         }
-    },
-    "lsp": {
-        "pyright": {
-            "settings": {
-                "python": {
-                    "pythonPath": ".hatch/fief/bin/python"
-                }
-            }
-        }
     }
 }

fief/_core.py

@@ -83,7 +83,9 @@ def __init__(self, fief: "Fief[U]", container: dishka.Container) -> None:
         self.container = container
 
     def get_storage(self, model: type[M]) -> StorageProtocol[M]:
-        component_key = self.fief._model_component_map.get(model)
+        component_key = self.fief._model_component_map.get(
+            model, dishka.DEFAULT_COMPONENT
+        )
         return self.container.get(StorageProtocol[model], component_key)  # type: ignore[valid-type]
 
     def get_method(self, type: type[MP], name: str | None = None) -> MP:
@@ -211,13 +213,19 @@ def close(self) -> None:
     def _init_storage(
         self, storage: StorageProvider | Sequence[StorageProvider]
     ) -> list["BaseProvider"]:
+        self._model_component_map = {}
+
         if isinstance(storage, StorageProvider):
-            storage = [storage]
+            return [storage]
 
         providers: list[BaseProvider] = []
         model_component_map: dict[type, str] = {}
         for i, s in enumerate(storage):
             component_key = f"storage_{i}"
+            if s.models is None:
+                raise ValueError(  # noqa: TRY003
+                    "You must provide models on the storage providers if you have multiple ones"
+                )
             for model in s.models:
                 if model in model_component_map:
                     raise ValueError(  # noqa: TRY003
@@ -240,7 +248,9 @@ def _init_methods(
                 raise ValueError(  # noqa: TRY003
                     f"Method with name {method.name} is already registered"
                 )
-            model_component_key = self._model_component_map[method.model]
+            model_component_key = self._model_component_map.get(
+                method.model, dishka.DEFAULT_COMPONENT
+            )
             provider = dishka.Provider(component=method.name)
             provider.provide(
                 method.get_provider(model_component_key), scope=dishka.Scope.REQUEST
@@ -259,7 +269,9 @@ def __init__(self, fief: "FiefAsync[U]", container: dishka.AsyncContainer) -> No
         self.container = container
 
     async def get_storage(self, model: type[M]) -> AsyncStorageProtocol[M]:
-        component_key = self.fief._model_component_map.get(model)
+        component_key = self.fief._model_component_map.get(
+            model, dishka.DEFAULT_COMPONENT
+        )
         return await self.container.get(AsyncStorageProtocol[model], component_key)  # type: ignore[valid-type]
 
     async def get_user_storage(self) -> AsyncStorageProtocol[U]:
@@ -389,13 +401,19 @@ async def close(self) -> None:
     def _init_storage(
         self, storage: AsyncStorageProvider | Sequence[AsyncStorageProvider]
     ) -> list["BaseProvider"]:
+        self._model_component_map = {}
+
         if isinstance(storage, AsyncStorageProvider):
-            storage = [storage]
+            return [storage]
 
         providers: list[BaseProvider] = []
         model_component_map: dict[type, str] = {}
         for i, s in enumerate(storage):
             component_key = f"storage_{i}"
+            if s.models is None:
+                raise ValueError(  # noqa: TRY003
+                    "You must provide models on the storage providers if you have multiple ones"
+                )
             for model in s.models:
                 if model in model_component_map:
                     raise ValueError(  # noqa: TRY003
@@ -420,7 +438,9 @@ def _init_methods(
                 raise ValueError(  # noqa: TRY003
                     f"Method with name {method.name} is already registered"
                 )
-            model_component_key = self._model_component_map[method.model]
+            model_component_key = self._model_component_map.get(
+                method.model, dishka.DEFAULT_COMPONENT
+            )
             provider = dishka.Provider(component=method.name)
             provider.provide(
                 method.get_provider(model_component_key), scope=dishka.Scope.REQUEST

fief/storage/_protocol.py

@@ -7,15 +7,14 @@
 
 
 class StorageProvider(dishka.Provider):
-    models: list[type]
+    models: list[type] | None
 
     def __init__(self, models: Sequence[type] | None = None):
         super().__init__()
-        self.models = list(models or [])
-        for model in self.models:
-            self.provide(self.get_model_provider(model), scope=dishka.Scope.REQUEST)
+        self.models = list(models) if models is not None else None
+        self.provide(self.get_provider(), scope=dishka.Scope.REQUEST)
 
-    def get_model_provider(self, model: type[M]) -> Callable[..., "StorageProtocol[M]"]:
+    def get_provider(self) -> Callable[..., "StorageProtocol[typing.Any]"]:
         raise NotImplementedError()
 
 
@@ -73,17 +72,14 @@ def delete(self, id: typing.Any) -> M | None:  # pragma: no cover
 
 
 class AsyncStorageProvider(dishka.Provider):
-    models: list[type]
+    models: list[type] | None
 
     def __init__(self, models: Sequence[type] | None = None):
         super().__init__()
-        self.models = list(models or [])
-        for model in self.models:
-            self.provide(self.get_model_provider(model), scope=dishka.Scope.REQUEST)
+        self.models = list(models) if models is not None else None
+        self.provide(self.get_provider(), scope=dishka.Scope.REQUEST)
 
-    def get_model_provider(
-        self, model: type[M]
-    ) -> Callable[..., "AsyncStorageProtocol[M]"]:
+    def get_provider(self) -> Callable[..., "AsyncStorageProtocol[typing.Any]"]:
         raise NotImplementedError()
 
 

fief/storage/sqlalchemy.py

@@ -76,10 +76,11 @@ def session(self, engine: Engine) -> Iterable[Session]:
                 raise
             session.commit()
 
-    def get_model_provider(self, model: type[M]) -> Callable[..., StorageProtocol[M]]:
+    def get_provider(self) -> Callable[..., StorageProtocol[typing.Any]]:
         def _provide(
+            model: type[M],
             session: Session,
-        ) -> StorageProtocol[model]:  # type: ignore[valid-type]
+        ) -> StorageProtocol[M]:
             return self.storage_class(model, session)
 
         return _provide
@@ -147,12 +148,11 @@ async def session(self, engine: AsyncEngine) -> AsyncIterable[AsyncSession]:
                 raise
             await session.commit()
 
-    def get_model_provider(
-        self, model: type[M]
-    ) -> Callable[..., AsyncStorageProtocol[M]]:
+    def get_provider(self) -> Callable[..., AsyncStorageProtocol[typing.Any]]:
         def _provide(
+            model: type[M],
             session: AsyncSession,
-        ) -> AsyncStorageProtocol[model]:  # type: ignore[valid-type]
+        ) -> AsyncStorageProtocol[M]:
             return self.storage_class(model, session)
 
         return _provide

tests/fixtures.py

@@ -81,8 +81,8 @@ async def delete(self, id: typing.Any) -> M | None:
 class MockProvider(StorageProvider):
     storage_class = MockStorage
 
-    def get_model_provider(self, model: type[M]) -> Callable[..., "StorageProtocol[M]"]:
-        def _provide() -> StorageProtocol[model]:  # type: ignore[valid-type]
+    def get_provider(self) -> Callable[..., StorageProtocol[typing.Any]]:
+        def _provide(model: type[M]) -> StorageProtocol[M]:
             return self.storage_class(model)
 
         return _provide
@@ -91,10 +91,8 @@ def _provide() -> StorageProtocol[model]:  # type: ignore[valid-type]
 class MockAsyncProvider(AsyncStorageProvider):
     storage_class = MockAsyncStorage
 
-    def get_model_provider(
-        self, model: type[M]
-    ) -> Callable[..., "AsyncStorageProtocol[M]"]:
-        def _provide() -> AsyncStorageProtocol[model]:  # type: ignore[valid-type]
+    def get_provider(self) -> Callable[..., AsyncStorageProtocol[typing.Any]]:
+        def _provide(model: type[M]) -> AsyncStorageProtocol[M]:
             return self.storage_class(model)
 
         return _provide

tests/test_core.py

@@ -53,6 +53,87 @@ class ModelB(Base):
             sa.Integer, primary_key=True, autoincrement=True
         )
 
+    fief = Fief(
+        storage=SQLAlchemyProvider("sqlite:///db.db"),
+        methods=(),
+        user_model=UserModel,
+    )
+
+    with fief as fief_request:
+        storage_a = fief_request.get_storage(ModelA)
+        assert isinstance(storage_a, SQLAlchemyStorage)
+        assert storage_a.model is ModelA
+        storage_a_bind = storage_a.session.bind
+        assert isinstance(storage_a_bind, sa.Engine)
+        assert str(storage_a_bind.url) == "sqlite:///db.db"
+
+        storage_b = fief_request.get_storage(ModelB)
+        assert isinstance(storage_b, SQLAlchemyStorage)
+        assert storage_b.model is ModelB
+        storage_b_bind = storage_b.session.bind
+        assert isinstance(storage_b_bind, sa.Engine)
+        assert str(storage_b_bind.url) == "sqlite:///db.db"
+
+    fief.close()
+
+
+@pytest.mark.parametrize("anyio_backend", ["asyncio"])
+async def test_async_storage_provider(anyio_backend: str) -> None:
+    class Base(DeclarativeBase):
+        pass
+
+    class ModelA(Base):
+        __tablename__ = "model_a"
+        id: Mapped[int] = mapped_column(
+            sa.Integer, primary_key=True, autoincrement=True
+        )
+
+    class ModelB(Base):
+        __tablename__ = "model_b"
+        id: Mapped[int] = mapped_column(
+            sa.Integer, primary_key=True, autoincrement=True
+        )
+
+    fief = FiefAsync(
+        storage=SQLAlchemyAsyncProvider("sqlite+aiosqlite:///db.db"),
+        methods=(),
+        user_model=UserModel,
+    )
+
+    async with fief as fief_request:
+        storage_a = await fief_request.get_storage(ModelA)
+        assert isinstance(storage_a, SQLAlchemyAsyncStorage)
+        assert storage_a.model is ModelA
+        storage_a_bind = storage_a.session.bind
+        assert isinstance(storage_a_bind, sa_asyncio.AsyncEngine)
+        assert str(storage_a_bind.url) == "sqlite+aiosqlite:///db.db"
+
+        storage_b = await fief_request.get_storage(ModelB)
+        assert isinstance(storage_b, SQLAlchemyAsyncStorage)
+        assert storage_b.model is ModelB
+        storage_b_bind = storage_b.session.bind
+        assert isinstance(storage_b_bind, sa_asyncio.AsyncEngine)
+        assert str(storage_b_bind.url) == "sqlite+aiosqlite:///db.db"
+
+    await fief.close()
+
+
+def test_multiple_storage_providers() -> None:
+    class Base(DeclarativeBase):
+        pass
+
+    class ModelA(Base):
+        __tablename__ = "model_a"
+        id: Mapped[int] = mapped_column(
+            sa.Integer, primary_key=True, autoincrement=True
+        )
+
+    class ModelB(Base):
+        __tablename__ = "model_b"
+        id: Mapped[int] = mapped_column(
+            sa.Integer, primary_key=True, autoincrement=True
+        )
+
     fief = Fief(
         storage=(
             SQLAlchemyProvider("sqlite:///db_a.db", models=[ModelA]),
@@ -86,7 +167,7 @@ class ModelB(Base):
 
 
 @pytest.mark.parametrize("anyio_backend", ["asyncio"])
-async def test_async_storage_provider(anyio_backend: str) -> None:
+async def test_async_multiple_storage_providers(anyio_backend: str) -> None:
     class Base(DeclarativeBase):
         pass
 
@@ -136,7 +217,7 @@ class ModelB(Base):
 
 def test_method_provider() -> None:
     fief = Fief(
-        storage=MockProvider(models=[UserModel, MethodModel]),
+        storage=MockProvider(),
         methods=(
             PasswordMethodProvider(MethodModel),
             PasswordMethodProvider(MethodModel, name="password2"),
@@ -156,7 +237,7 @@ def test_method_provider() -> None:
 @pytest.mark.anyio
 async def test_async_method_provider() -> None:
     fief = FiefAsync(
-        storage=MockAsyncProvider(models=[UserModel, MethodModel]),
+        storage=MockAsyncProvider(),
         methods=(
             PasswordAsyncMethodProvider(MethodModel),
             PasswordAsyncMethodProvider(MethodModel, name="password2"),
@@ -186,10 +267,7 @@ class PasswordMethodModel:
 @pytest.fixture
 def fief() -> Fief[UserModel]:
     return Fief(
-        storage=(
-            MockProvider(models=[UserModel]),
-            MockProvider(models=[PasswordMethodModel]),
-        ),
+        storage=MockProvider(),
         methods=(PasswordMethodProvider(PasswordMethodModel),),
         user_model=UserModel,
     )
@@ -204,10 +282,7 @@ def fief_request(fief: Fief[UserModel]) -> Generator[FiefRequest[UserModel]]:
 @pytest.fixture
 def fief_async() -> FiefAsync[UserModel]:
     return FiefAsync(
-        storage=(
-            MockAsyncProvider(models=[UserModel]),
-            MockAsyncProvider(models=[PasswordMethodModel]),
-        ),
+        storage=MockAsyncProvider(),
         methods=(PasswordAsyncMethodProvider(PasswordMethodModel),),
         user_model=UserModel,
     )