Merge pull request #26 from frankie567/issue-20-hooks

Fix #20: revamp and define new event handlers.
2026-03-13 07:49:55 +08:00 · 2019-10-25 08:33:23 +02:00
parent af12834225 3506df31e7
commit 8987b1bb2a
6 changed files with 165 additions and 94 deletions
--- a/docs/configuration/router.md
+++ b/docs/configuration/router.md
@@ -2,33 +2,13 @@

 We're almost there! The last step is to configure the `FastAPIUsers` object that will wire the database adapter, the authentication class and the user model to expose the FastAPI router.

-## Hooks
-
-In order to be as unopinionated as possible, you'll have to define your logic after some actions.
-
-### After forgot password
-
-This hook is called after a successful forgot password request. It is called with **two arguments**: the **user** which has requested to reset their password and a ready-to-use **JWT token** that will be accepted by the reset password route.
-
-Typically, you'll want to **send an e-mail** with the link (and the token) that allows the user to reset their password.
-
-You can define it as an `async` or standard method.
-
-Example:
-
-```py
-def on_after_forgot_password(user, token):
-    print(f'User {user.id} has forgot their password. Reset token: {token}')
-```
-
 ## Configure `FastAPIUsers`

-The last step is to instantiate `FastAPIUsers` object with all the elements we defined before. More precisely:
+Configure `FastAPIUsers` object with all the elements we defined before. More precisely:

 * `db`: Database adapter instance.
 * `auth`: Authentication logic instance.
 * `user_model`: Pydantic model of a user.
-* `on_after_forgot_password`: Hook called after a forgot password request.
 * `reset_password_token_secret`: Secret to encode reset password token.
 * `reset_password_token_lifetime_seconds`: Lifetime of reset password token in seconds. Default to one hour.

@@ -39,7 +19,6 @@ fastapi_users = FastAPIUsers(
    user_db,
    auth,
    User,
-    on_after_forgot_password,
    SECRET,
 )
 ```
@@ -51,6 +30,42 @@ app = FastAPI()
 app.include_router(fastapi_users.router, prefix="/users", tags=["users"])
 ```

+## Event handlers
+
+In order to be as unopinionated as possible, we expose decorators that allow you to plug your own logic after some actions. You can have several handlers per event.
+
+### After register
+
+This event handler is called after a successful registration. It is called with **one argument**: the **user** that has just registered.
+
+Typically, you'll want to **send a welcome e-mail** or add it to your marketing analytics pipeline.
+
+You can define it as an `async` or standard method.
+
+Example:
+
+```py
+@fastapi_users.on_after_register()
+def on_after_register(user: User):
+    print(f"User {user.id} has registered.")
+```
+
+### After forgot password
+
+This event handler is called after a successful forgot password request. It is called with **two arguments**: the **user** which has requested to reset their password and a ready-to-use **JWT token** that will be accepted by the reset password route.
+
+Typically, you'll want to **send an e-mail** with the link (and the token) that allows the user to reset their password.
+
+You can define it as an `async` or standard method.
+
+Example:
+
+```py
+@fastapi_users.on_after_forgot_password()
+def on_after_forgot_password(user: User, token: str):
+    print(f"User {user.id} has forgot their password. Reset token: {token}")
+```
+
 ## Next steps

 Check out a [full example](full_example.md) that will show you the big picture.
--- a/docs/src/full_sqlalchemy.py
+++ b/docs/src/full_sqlalchemy.py
@@ -35,16 +35,21 @@ class User(BaseUser):

 auth = JWTAuthentication(secret=SECRET, lifetime_seconds=3600)

-
-def on_after_forgot_password(user, token):
-    print(f"User {user.id} has forgot their password. Reset token: {token}")
-
-
 app = FastAPI()
-fastapi_users = FastAPIUsers(user_db, auth, User, on_after_forgot_password, SECRET)
+fastapi_users = FastAPIUsers(user_db, auth, User, SECRET)
 app.include_router(fastapi_users.router, prefix="/users", tags=["users"])


+@fastapi_users.on_after_register()
+def on_after_register(user: User):
+    print(f"User {user.id} has registered.")
+
+
+@fastapi_users.on_after_forgot_password()
+def on_after_forgot_password(user: User, token: str):
+    print(f"User {user.id} has forgot their password. Reset token: {token}")
+
+
@app.on_event("startup")
 async def startup():
    await database.connect()
--- a/fastapi_users/fastapi_users.py
+++ b/fastapi_users/fastapi_users.py
@@ -1,11 +1,9 @@
-from typing import Any, Callable, Type
-
-from fastapi import APIRouter
+from typing import Callable, Type

 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
 from fastapi_users.models import BaseUser, BaseUserDB
-from fastapi_users.router import get_user_router
+from fastapi_users.router import Events, UserRouter, get_user_router


 class FastAPIUsers:
@@ -15,17 +13,16 @@ class FastAPIUsers:
    :param db: Database adapter instance.
    :param auth: Authentication logic instance.
    :param user_model: Pydantic model of a user.
-    :param on_after_forgot_password: Hook called after a forgot password request.
    :param reset_password_token_secret: Secret to encode reset password token.
    :param reset_password_token_lifetime_seconds: Lifetime of reset password token.

-    :attribute router: FastAPI router exposing authentication routes.
+    :attribute router: Router exposing authentication routes.
    :attribute get_current_user: Dependency callable to inject authenticated user.
    """

    db: BaseUserDatabase
    auth: BaseAuthentication
-    router: APIRouter
+    router: UserRouter
    get_current_user: Callable[..., BaseUserDB]

    def __init__(
@@ -33,7 +30,6 @@ class FastAPIUsers:
        db: BaseUserDatabase,
        auth: BaseAuthentication,
        user_model: Type[BaseUser],
-        on_after_forgot_password: Callable[[BaseUserDB, str], Any],
        reset_password_token_secret: str,
        reset_password_token_lifetime_seconds: int = 3600,
    ):
@@ -43,7 +39,6 @@ class FastAPIUsers:
            self.db,
            user_model,
            self.auth,
-            on_after_forgot_password,
            reset_password_token_secret,
            reset_password_token_lifetime_seconds,
        )
@@ -56,3 +51,18 @@ class FastAPIUsers:

        get_current_superuser = self.auth.get_current_superuser(self.db)
        self.get_current_superuser = get_current_superuser  # type: ignore
+
+    def on_after_register(self) -> Callable:
+        """Add an event handler on successful registration."""
+        return self._on_event(Events.ON_AFTER_REGISTER)
+
+    def on_after_forgot_password(self) -> Callable:
+        """Add an event handler on successful forgot password request."""
+        return self._on_event(Events.ON_AFTER_FORGOT_PASSWORD)
+
+    def _on_event(self, event_type: Events) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            self.router.add_event_handler(event_type, func)
+            return func
+
+        return decorator
--- a/fastapi_users/router.py
+++ b/fastapi_users/router.py
@@ -1,5 +1,7 @@
 import asyncio
-from typing import Any, Callable, Type
+import typing
+from collections import defaultdict
+from enum import Enum

 import jwt
 from fastapi import APIRouter, Body, Depends, HTTPException
@@ -10,27 +12,46 @@ from starlette.responses import Response

 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
-from fastapi_users.models import BaseUser, BaseUserDB, Models
+from fastapi_users.models import BaseUser, Models
 from fastapi_users.password import get_password_hash
 from fastapi_users.utils import JWT_ALGORITHM, generate_jwt


+class Events(Enum):
+    ON_AFTER_REGISTER = 1
+    ON_AFTER_FORGOT_PASSWORD = 2
+
+
+class UserRouter(APIRouter):
+    event_handlers: typing.DefaultDict[Events, typing.List[typing.Callable]]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.event_handlers = defaultdict(list)
+
+    def add_event_handler(self, event_type: Events, func: typing.Callable) -> None:
+        self.event_handlers[event_type].append(func)
+
+    async def run_handlers(self, event_type: Events, *args, **kwargs) -> None:
+        for handler in self.event_handlers[event_type]:
+            if asyncio.iscoroutinefunction(handler):
+                await handler(*args, **kwargs)
+            else:
+                handler(*args, **kwargs)
+
+
 def get_user_router(
    user_db: BaseUserDatabase,
-    user_model: Type[BaseUser],
+    user_model: typing.Type[BaseUser],
    auth: BaseAuthentication,
-    on_after_forgot_password: Callable[[BaseUserDB, str], Any],
    reset_password_token_secret: str,
    reset_password_token_lifetime_seconds: int = 3600,
-) -> APIRouter:
+) -> UserRouter:
    """Generate a router with the authentication routes."""
-    router = APIRouter()
+    router = UserRouter()
    models = Models(user_model)

    reset_password_token_audience = "fastapi-users:reset"
-    is_on_after_forgot_password_async = asyncio.iscoroutinefunction(
-        on_after_forgot_password
-    )

    get_current_active_user = auth.get_current_active_user(user_db)

@@ -48,6 +69,9 @@ def get_user_router(
            **user.create_update_dict(), hashed_password=hashed_password
        )
        created_user = await user_db.create(db_user)
+
+        await router.run_handlers(Events.ON_AFTER_REGISTER, created_user)
+
        return created_user

    @router.post("/login")
@@ -74,10 +98,7 @@ def get_user_router(
                reset_password_token_lifetime_seconds,
                reset_password_token_secret,
            )
-            if is_on_after_forgot_password_async:
-                await on_after_forgot_password(user, token)
-            else:
-                on_after_forgot_password(user, token)
+            await router.run_handlers(Events.ON_AFTER_FORGOT_PASSWORD, user, token)

        return None

--- a/tests/test_fastapi_users.py
+++ b/tests/test_fastapi_users.py
@@ -5,26 +5,37 @@ from starlette.testclient import TestClient

 from fastapi_users import FastAPIUsers
 from fastapi_users.models import BaseUser, BaseUserDB
-
-SECRET = "SECRET"
+from fastapi_users.router import Events


-def sync_on_after_forgot_password():
+def sync_event_handler():
    return None


-async def async_on_after_forgot_password():
+async def async_event_handler():
    return None


-@pytest.fixture(params=[sync_on_after_forgot_password, async_on_after_forgot_password])
-def test_app_client(request, mock_user_db, mock_authentication) -> TestClient:
+@pytest.fixture(params=[sync_event_handler, async_event_handler])
+def fastapi_users(request, mock_user_db, mock_authentication) -> FastAPIUsers:
    class User(BaseUser):
        pass

-    fastapi_users = FastAPIUsers(
-        mock_user_db, mock_authentication, User, request.param, SECRET
-    )
+    fastapi_users = FastAPIUsers(mock_user_db, mock_authentication, User, "SECRET")
+
+    @fastapi_users.on_after_register()
+    def on_after_register():
+        return request.param()
+
+    @fastapi_users.on_after_forgot_password()
+    def on_after_forgot_password():
+        return request.param()
+
+    return fastapi_users
+
+
+@pytest.fixture()
+def test_app_client(fastapi_users) -> TestClient:
    app = FastAPI()
    app.include_router(fastapi_users.router)

@@ -43,6 +54,13 @@ def test_app_client(request, mock_user_db, mock_authentication) -> TestClient:
    return TestClient(app)


+class TestFastAPIUsers:
+    def test_event_handlers(self, fastapi_users):
+        event_handlers = fastapi_users.router.event_handlers
+        assert len(event_handlers[Events.ON_AFTER_REGISTER]) == 1
+        assert len(event_handlers[Events.ON_AFTER_FORGOT_PASSWORD]) == 1
+
+
 class TestRouter:
    def test_routes_exist(self, test_app_client: TestClient):
        response = test_app_client.post("/register")
--- a/tests/test_router.py
+++ b/tests/test_router.py
@@ -8,7 +8,7 @@ from starlette import status
 from starlette.testclient import TestClient

 from fastapi_users.models import BaseUser, BaseUserDB
-from fastapi_users.router import get_user_router
+from fastapi_users.router import Events, get_user_router
 from fastapi_users.utils import JWT_ALGORITHM, generate_jwt

 SECRET = "SECRET"
@@ -26,35 +26,31 @@ def forgot_password_token():
    return _forgot_password_token


-def on_after_forgot_password_sync():
+def event_handler_sync():
    return MagicMock(return_value=None)


-def on_after_forgot_password_async():
+def event_handler_async():
    return asynctest.CoroutineMock(return_value=None)


-@pytest.fixture(params=[on_after_forgot_password_sync, on_after_forgot_password_async])
-def on_after_forgot_password(request):
+@pytest.fixture(params=[event_handler_sync, event_handler_async])
+def event_handler(request):
    return request.param()


@pytest.fixture()
-def test_app_client(
-    mock_user_db, mock_authentication, on_after_forgot_password
-) -> TestClient:
+def test_app_client(mock_user_db, mock_authentication, event_handler) -> TestClient:
    class User(BaseUser):
        pass

    userRouter = get_user_router(
-        mock_user_db,
-        User,
-        mock_authentication,
-        on_after_forgot_password,
-        SECRET,
-        LIFETIME,
+        mock_user_db, User, mock_authentication, SECRET, LIFETIME
    )

+    userRouter.add_event_handler(Events.ON_AFTER_REGISTER, event_handler)
+    userRouter.add_event_handler(Events.ON_AFTER_FORGOT_PASSWORD, event_handler)
+
    app = FastAPI()
    app.include_router(userRouter)

@@ -62,36 +58,44 @@ def test_app_client(


 class TestRegister:
-    def test_empty_body(self, test_app_client: TestClient):
+    def test_empty_body(self, test_app_client: TestClient, event_handler):
        response = test_app_client.post("/register", json={})
        assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+        assert event_handler.called is False

-    def test_missing_password(self, test_app_client: TestClient):
+    def test_missing_password(self, test_app_client: TestClient, event_handler):
        json = {"email": "king.arthur@camelot.bt"}
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+        assert event_handler.called is False

-    def test_wrong_email(self, test_app_client: TestClient):
+    def test_wrong_email(self, test_app_client: TestClient, event_handler):
        json = {"email": "king.arthur", "password": "guinevere"}
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+        assert event_handler.called is False

-    def test_existing_user(self, test_app_client: TestClient):
+    def test_existing_user(self, test_app_client: TestClient, event_handler):
        json = {"email": "king.arthur@camelot.bt", "password": "guinevere"}
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_400_BAD_REQUEST
+        assert event_handler.called is False

-    def test_valid_body(self, test_app_client: TestClient):
+    def test_valid_body(self, test_app_client: TestClient, event_handler):
        json = {"email": "lancelot@camelot.bt", "password": "guinevere"}
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_201_CREATED
+        assert event_handler.called is True

        response_json = response.json()
        assert "hashed_password" not in response_json
        assert "password" not in response_json
        assert response_json["id"] is not None

-    def test_valid_body_is_superuser(self, test_app_client: TestClient):
+        actual_user = event_handler.call_args[0][0]
+        assert actual_user.id == response_json["id"]
+
+    def test_valid_body_is_superuser(self, test_app_client: TestClient, event_handler):
        json = {
            "email": "lancelot@camelot.bt",
            "password": "guinevere",
@@ -99,11 +103,12 @@ class TestRegister:
        }
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_201_CREATED
+        assert event_handler.called is True

        response_json = response.json()
        assert response_json["is_superuser"] is False

-    def test_valid_body_is_active(self, test_app_client: TestClient):
+    def test_valid_body_is_active(self, test_app_client: TestClient, event_handler):
        json = {
            "email": "lancelot@camelot.bt",
            "password": "guinevere",
@@ -111,6 +116,7 @@ class TestRegister:
        }
        response = test_app_client.post("/register", json=json)
        assert response.status_code == status.HTTP_201_CREATED
+        assert event_handler.called is True

        response_json = response.json()
        assert response_json["is_active"] is True
@@ -154,36 +160,32 @@ class TestLogin:


 class TestForgotPassword:
-    def test_empty_body(self, test_app_client: TestClient, on_after_forgot_password):
+    def test_empty_body(self, test_app_client: TestClient, event_handler):
        response = test_app_client.post("/forgot-password", json={})
        assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
-        assert on_after_forgot_password.called is False
+        assert event_handler.called is False

-    def test_not_existing_user(
-        self, test_app_client: TestClient, on_after_forgot_password
-    ):
+    def test_not_existing_user(self, test_app_client: TestClient, event_handler):
        json = {"email": "lancelot@camelot.bt"}
        response = test_app_client.post("/forgot-password", json=json)
        assert response.status_code == status.HTTP_202_ACCEPTED
-        assert on_after_forgot_password.called is False
+        assert event_handler.called is False

-    def test_inactive_user(self, test_app_client: TestClient, on_after_forgot_password):
+    def test_inactive_user(self, test_app_client: TestClient, event_handler):
        json = {"email": "percival@camelot.bt"}
        response = test_app_client.post("/forgot-password", json=json)
        assert response.status_code == status.HTTP_202_ACCEPTED
-        assert on_after_forgot_password.called is False
+        assert event_handler.called is False

-    def test_existing_user(
-        self, test_app_client: TestClient, on_after_forgot_password, user
-    ):
+    def test_existing_user(self, test_app_client: TestClient, event_handler, user):
        json = {"email": "king.arthur@camelot.bt"}
        response = test_app_client.post("/forgot-password", json=json)
        assert response.status_code == status.HTTP_202_ACCEPTED
-        assert on_after_forgot_password.called is True
+        assert event_handler.called is True

-        actual_user = on_after_forgot_password.call_args[0][0]
+        actual_user = event_handler.call_args[0][0]
        assert actual_user.id == user.id
-        actual_token = on_after_forgot_password.call_args[0][1]
+        actual_token = event_handler.call_args[0][1]
        decoded_token = jwt.decode(
            actual_token,
            SECRET,