From 88260e371480204ab33f9103e0dcc3b3b27d3a4e Mon Sep 17 00:00:00 2001
From: Yasser Tahiri <yasserth19@gmail.com>
Date: Thu, 4 Apr 2024 03:14:28 +0000
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20FastAPI=20dependencies=20for?=
 =?UTF-8?q?=20token=20operations=20in=20route=20logic=20(#566)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 🔧 Add methods for setting subject getter and token blocklist callbacks

* ✨ Add FastAPI dependencies for token operations in route logic
---
 authx/_internal/_callback.py |  8 ++++++
 authx/main.py                | 50 +++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/authx/_internal/_callback.py b/authx/_internal/_callback.py
index 5e6f9dbe..0337f116 100644
--- a/authx/_internal/_callback.py
+++ b/authx/_internal/_callback.py
@@ -66,6 +66,14 @@ def set_callback_token_blocklist(self, callback: TokenCallback) -> None:
         """Set callback for token"""
         self.callback_is_token_in_blocklist = callback
 
+    def set_subject_getter(self, callback: ModelCallback[T]) -> None:
+        """Set the callback to run for subject retrieval and serialization"""
+        self.set_callback_get_model_instance(callback)
+
+    def set_token_blocklist(self, callback: TokenCallback) -> None:
+        """Set the callback to run for validation of revoked tokens"""
+        self.set_callback_token_blocklist(callback)
+
     def _get_current_subject(self, uid: str, **kwargs) -> T:
         """Get current model instance from callback"""
         self._check_model_callback_is_set()
diff --git a/authx/main.py b/authx/main.py
index 3d44f4aa..2ec550b7 100644
--- a/authx/main.py
+++ b/authx/main.py
@@ -1,7 +1,7 @@
 import contextlib
 from typing import Any, Callable, Coroutine, Dict, Literal, Optional, overload
 
-from fastapi import Request, Response
+from fastapi import Depends, Request, Response
 
 from authx._internal._callback import _CallbackHandler
 from authx._internal._error import _ErrorHandler
@@ -501,6 +501,54 @@ def unset_cookies(self, response: Response) -> None:
         self.unset_access_cookies(response)
         self.unset_refresh_cookies(response)
 
+    # Notes:
+    # The AuthXDeps is a utility class, to enable quick token operations
+    # within the route logic. It provides methods to avoid addtional code
+    # in your route that would be outside of the route logic
+
+    # Such methods includes setting and unsetting cookies without the need
+    # to generate a response object beforhand.
+
+    @property
+    def DEPENDENCY(self) -> AuthXDependency:
+        """FastAPI Dependency to return an AuthX sub-object within the route context"""
+        return Depends(self.get_dependency)
+
+    @property
+    def BUNDLE(self) -> AuthXDependency:
+        """FastAPI Dependency to return a AuthX sub-object within the route context"""
+        return self.DEPENDENCY
+
+    @property
+    def FRESH_REQUIRED(self) -> TokenPayload:
+        """FastAPI Dependency to enforce valid token availability in request"""
+        return Depends(self.fresh_token_required)
+
+    @property
+    def ACCESS_REQUIRED(self) -> TokenPayload:
+        """FastAPI Dependency to enforce presence of an `access` token in request"""
+        return Depends(self.access_token_required)
+
+    @property
+    def REFRESH_REQUIRED(self) -> TokenPayload:
+        """FastAPI Dependency to enforce presence of a `refresh` token in request"""
+        return Depends(self.refresh_token_required)
+
+    @property
+    def ACCESS_TOKEN(self) -> RequestToken:
+        """FastAPI Dependency to retrieve access token from request"""
+        return Depends(self.get_token_from_request(type="access"))
+
+    @property
+    def REFRESH_TOKEN(self) -> RequestToken:
+        """FastAPI Dependency to retrieve refresh token from request"""
+        return Depends(self.get_token_from_request(type="refresh"))
+
+    @property
+    def CURRENT_SUBJECT(self) -> T:
+        """FastAPI Dependency to retrieve the current subject from request"""
+        return Depends(self.get_current_subject)
+
     def get_dependency(self, request: Request, response: Response) -> AuthXDependency:
         """FastAPI Dependency to return a AuthX sub-object within the route context