Improving user management and auto-refreshing user firms

api/rpk-api/hub/auth/__init__.py

@@ -1,19 +1,19 @@
 import os
-from typing import Any, Optional
+from typing import Any
 
-from beanie import PydanticObjectId
+from beanie import PydanticObjectId, Document
 from fastapi import Depends, Response, status
 from fastapi_users import BaseUserManager, FastAPIUsers, schemas, models
-from fastapi_users.authentication import AuthenticationBackend, BearerTransport, CookieTransport, Strategy
+from fastapi_users.authentication import AuthenticationBackend, CookieTransport, Strategy
 from fastapi_users.authentication.strategy import AccessTokenDatabase, DatabaseStrategy
-from fastapi_users_db_beanie.access_token import BeanieBaseAccessTokenDocument, BeanieAccessTokenDatabase
+from fastapi_users_db_beanie.access_token import BeanieBaseAccessToken, BeanieAccessTokenDatabase
 from fastapi_users.openapi import OpenAPIResponseType
 from httpx_oauth.clients.google import GoogleOAuth2
 from httpx_oauth.clients.discord import DiscordOAuth2
 from starlette.responses import JSONResponse, RedirectResponse
 
 from hub.user import User, get_user_db
-from hub.user.schemas import UserSchema
+from hub.user.schemas import UserSchema, UserUpdateSchema
 
 
 SECRET = os.getenv("FASTAPI_USERS_SECRET")
@@ -23,7 +23,7 @@ discord_oauth_client = DiscordOAuth2(os.getenv("DISCORD_CLIENT_ID"), os.getenv("
 TOKEN_LIFETIME = 3600
 
 
-class AccessToken(BeanieBaseAccessTokenDocument):
+class AccessToken(BeanieBaseAccessToken, Document):
     pass
 
 async def get_access_token_db():
@@ -84,10 +84,11 @@ auth_router = fastapi_users.get_auth_router(auth_backend, requires_verification=
 register_router = fastapi_users.get_register_router(UserSchema, schemas.BaseUserCreate)
 password_router = fastapi_users.get_reset_password_router()
 verification_router = fastapi_users.get_verify_router(UserSchema)
-users_router = fastapi_users.get_users_router(UserSchema, schemas.BaseUserUpdate)
+users_router = fastapi_users.get_users_router(UserSchema, UserUpdateSchema)
 
 cookie_transport = CookieTransportOauth(cookie_name="rpkapiusersauth")
 auth_backend = AuthenticationBackend(name="db", transport=cookie_transport, get_strategy=get_database_strategy, )
+
 google_oauth_router = fastapi_users.get_oauth_router(google_oauth_client, auth_backend, SECRET, is_verified_by_default=True)
 discord_oauth_router = fastapi_users.get_oauth_router(discord_oauth_client, auth_backend, SECRET, is_verified_by_default=True)
 

api/rpk-api/hub/db.py

@@ -18,4 +18,4 @@ async def init_db():
 
     await init_beanie(database=client.hub,
                       document_models=[User, AccessToken, Firm],
-                      allow_index_dropping=True)
+                      allow_index_dropping=True)

api/rpk-api/hub/firm/routes.py

@@ -4,44 +4,41 @@ from fastapi import APIRouter, Depends, HTTPException
 from hub.auth import get_current_user
 from hub.firm import Firm, FirmRead, FirmCreate, FirmUpdate
 
-model = Firm
-model_read = FirmRead
-model_create = FirmCreate
-model_update = FirmUpdate
 
 router = APIRouter()
 
-@router.post("/", response_description="{} added to the database".format(model.__name__))
-async def create(item: model_create, user=Depends(get_current_user)) -> model_read:
-    exists = await Firm.find_one({"name": item.name, "instance": item.instance})
+@router.post("/", response_description="{} added to the database".format(Firm.__name__))
+async def create(item: FirmCreate, user=Depends(get_current_user)) -> FirmRead:
+    firm_dict = {"name": item.name, "instance": item.instance}
+    exists = await Firm.find_one(firm_dict)
     if exists:
         raise HTTPException(status_code=400, detail="Firm already exists")
 
-    record = model(created_by=user.id, updated_by=user.id, owner=user.id, **item.model_dump())
+    record = Firm(created_by=user.id, updated_by=user.id, owner=user.id, **item.model_dump())
     o = await record.create()
-    user.firms.append(o.id)
-    user.save()
-    return model_read(**o.model_dump())
+    user.firms.append(firm_dict)
+    await user.save()
+    return FirmRead(**o.model_dump())
 
-@router.get("/{id}", response_description="{} record retrieved".format(model.__name__))
-async def read_id(id: PydanticObjectId, user=Depends(get_current_user)) -> model_read:
-    item = await model.get(id)
-    return model_read(**item.model_dump())
+@router.get("/{id}", response_description="{} record retrieved".format(Firm.__name__))
+async def read_id(id: PydanticObjectId, user=Depends(get_current_user)) -> FirmRead:
+    item = await Firm.get(id)
+    return FirmRead(**item.model_dump())
 
 
-@router.put("/{id}", response_description="{} record updated".format(model.__name__))
-async def update(id: PydanticObjectId, req: model_update, user=Depends(get_current_user)) -> model_read:
-    item = await model.get(id)
+@router.put("/{id}", response_description="{} record updated".format(Firm.__name__))
+async def update(id: PydanticObjectId, req: FirmUpdate, user=Depends(get_current_user)) -> FirmRead:
+    item = await Firm.get(id)
     if not item:
         raise HTTPException(
             status_code=404,
-            detail="{} record not found!".format(model.__name__)
+            detail="{} record not found!".format(Firm.__name__)
         )
 
     if item.owner != user.id:
         raise HTTPException(
             status_code=403,
-            detail="Insufficient credentials to modify {} record".format(model.__name__)
+            detail="Insufficient credentials to modify {} record".format(Firm.__name__)
         )
 
     req = {k: v for k, v in req.model_dump().items() if v is not None}
@@ -50,23 +47,23 @@ async def update(id: PydanticObjectId, req: model_update, user=Depends(get_curre
     }}
 
     await item.update(update_query)
-    return model_read(**item.dict())
+    return FirmRead(**item.dict())
 
-@router.delete("/{id}", response_description="{} record deleted from the database".format(model.__name__))
+@router.delete("/{id}", response_description="{} record deleted from the database".format(Firm.__name__))
 async def delete(id: PydanticObjectId, user=Depends(get_current_user)) -> dict:
-    item = await model.get(id)
+    item = await Firm.get(id)
     if not item:
         raise HTTPException(
             status_code=404,
-            detail="{} record not found!".format(model.__name__)
+            detail="{} record not found!".format(Firm.__name__)
         )
     if item.owner != user.id:
         raise HTTPException(
             status_code=403,
-            detail="Insufficient credentials delete {} record".format(model.__name__)
+            detail="Insufficient credentials delete {} record".format(Firm.__name__)
         )
 
     await item.delete()
     return {
-        "message": "{} deleted successfully".format(model.__name__)
+        "message": "{} deleted successfully".format(Firm.__name__)
     }

api/rpk-api/hub/user/__init__.py

@@ -1,13 +1,17 @@
-from beanie import PydanticObjectId
-from fastapi_users_db_beanie import BaseOAuthAccount, BeanieUserDatabase, BeanieBaseUserDocument
+from beanie import Document
+from fastapi_users_db_beanie import BaseOAuthAccount, BeanieUserDatabase, BeanieBaseUser
 from pydantic import Field
 
+from hub.firm import FirmRead
+from hub.user.schemas import UserSchema, UserUpdateSchema
+
+
 class OAuthAccount(BaseOAuthAccount):
     pass
 
-class User(BeanieBaseUserDocument):
+class User(BeanieBaseUser, Document):
     oauth_accounts: list[OAuthAccount] = Field(default_factory=list)
-    firms: list[PydanticObjectId] = Field(default_factory=list)
+    firms: list[FirmRead] = Field(default_factory=list)
 
 class UserDatabase(BeanieUserDatabase):
     pass

api/rpk-api/hub/user/schemas.py

@@ -1,7 +1,12 @@
 from beanie import PydanticObjectId
-from fastapi_users.schemas import BaseUser
+from fastapi_users.schemas import BaseUser, BaseUserUpdate
 from pydantic import Field
 
+from hub.firm import FirmRead
+
 
 class UserSchema(BaseUser[PydanticObjectId]):
-    firms: list[PydanticObjectId] = Field()
+    firms: list[FirmRead] = Field()
+
+class UserUpdateSchema(BaseUserUpdate):
+    pass