abstracted out piccolo functionality into token provider

piccolo_api/token_auth/endpoints.py

@@ -1,3 +1,6 @@
+from abc import abstractmethod
+import typing as t
+
 from starlette.endpoints import HTTPEndpoint
 from starlette.exceptions import HTTPException
 from starlette.responses import JSONResponse
@@ -7,7 +10,32 @@
 from .tables import TokenAuth
 
 
+class TokenProvider:
+    """
+    Subclass this to provide your own custom token provider.
+    """
+
+    @abstractmethod
+    async def get_token(self, username: str, password: str) -> t.Optional[str]:
+        pass
+
+
+class PiccoloTokenProvider(TokenProvider):
+    async def get_token(self, username: str, password: str) -> t.Optional[str]:
+        user = await BaseUser.login(username=username, password=password)
+        if user:
+            return (
+                TokenAuth.select(TokenAuth.token)
+                .first()
+                .where(TokenAuth.user == user)
+            )
+        return None
+
+
 class TokenAuthLoginEndpoint(HTTPEndpoint):
+
+    token_provider: TokenProvider = PiccoloTokenProvider()
+
     async def post(self, request: Request) -> JSONResponse:
         """
         Return a token if the credentials are correct.
@@ -16,13 +44,10 @@ async def post(self, request: Request) -> JSONResponse:
         username = json.get("username")
         password = json.get("password")
         if username and password:
-            user = await BaseUser.login(username=username, password=password)
-            if user:
-                token = (
-                    TokenAuth.select(TokenAuth.token)
-                    .first()
-                    .where(TokenAuth.user == user)
-                )
+            token = await self.token_provider.get_token(
+                username=username, password=password
+            )
+            if token:
                 return JSONResponse({"token": token})
             else:
                 raise HTTPException(