add create notification endpoint + refactor jwt token validation

Co-authored-by: violeta <50527601+violetaperezandrade@users.noreply.github.com>

Author: fjpacheco

app/controller/Users.py

@@ -64,10 +64,17 @@ async def handle_login(self, auth_code: str):
             headers={"x-access-token": f"{jwt_token}"},
         )
 
-    def handle_update_user(self, update_data: dict, request: Request):
-        user_id = self.users_service.retrieve_user_id(request)
+    def handle_update_user(self, update_data: dict, user_id: int):
         self.users_service.update_user(user_id, update_data)
         return {
             "message": "User updated successfully",
             "status": status.HTTP_200_OK,
         }
+
+    def handle_create_notification(self, notification_data: dict,
+                                   user_id: int):
+        self.users_service.create_notification(user_id, notification_data)
+        return {
+            "message": "Notification created successfully",
+            "status": status.HTTP_201_CREATED,
+        }

app/docker/tablas.sql

@@ -29,21 +29,22 @@ VALUES ('Agus', 'agus@fi.uba.ar', TO_DATE('1999-01-29', 'YYYY-MM-DD'), '{"lat":
 --     ('Sofi', 'sofi@fi.uba.ar', TO_DATE('1998-04-26', 'YYYY-MM-DD'), '{"lat": 1190, "long": 500}', 'QWEQWE'),
 --     ('Violeta', 'violeta@fi.uba.ar', TO_DATE('1998-05-12', 'YYYY-MM-DD'), '{"lat": 330, "long": 2333}', 'ZXZXZX');
 
+DROP TABLE IF EXISTS users_service.alarms CASCADE;
+
 CREATE TABLE IF NOT EXISTS users_service.alarms (
     id SERIAL PRIMARY KEY,
     id_user INT NOT NULL REFERENCES users_service.users(id) ON DELETE CASCADE,
-    datetime TIMESTAMP WITH TIME ZONE NOT NULL,
+    date_time TIMESTAMP WITH TIME ZONE NOT NULL,
     content VARCHAR(128) NOT NULL
 );
-
-CREATE INDEX idx_alarms_datetime ON users_service.alarms(datetime);
+CREATE INDEX idx_alarms_date_time ON users_service.alarms(date_time);
 
 INSERT INTO
-    users_service.alarms (id_user, datetime, content)
-VALUES (1, '2024-04-29 02:15:00-03', 'Wake up 1!'),
-    (2, '2024-04-29 02:15:00-03', 'Wake up 2!'),
-    (3, '2024-04-29 02:25:00-03', 'Wake up 3!'),
-    (4, '2024-04-29 02:25:00-03', 'Wake up 4!');
+    users_service.alarms (id_user, date_time, content)
+VALUES (1, '2024-04-30 00:50:00-03', 'Wake up 1!'),
+    (1, '2024-04-30 00:50:00-03', 'Wake up 2!'),
+    (1, '2024-04-30 00:50:00-03', 'Wake up 3!'),
+    (1, '2024-04-30 00:50:00-03', 'Wake up 4!');
 
 -- SELECT users_service.alarms.id, users_service.alarms.content, users_service.users.device_token
 -- FROM users_service.alarms 

app/main.py

@@ -1,10 +1,17 @@
-from fastapi import FastAPI, Request, Query
+from fastapi import Depends, FastAPI, Query
 from controller.Users import UsersController
 from service.Users import UsersService
 from repository.Users import UsersRepository
-from schemas.Schemas import CreateUserSchema, UpdateUserSchema
-from schemas.Schemas import LoginRequest
 from typing import List, Annotated, Union
+from schemas.Schemas import (
+    CreateUserSchema,
+    UpdateUserSchema,
+    LoginRequest,
+    CreateNotificationSchema
+)
+from security.JWTBearer import get_current_user_id
+from logs import init_logging
+logger = init_logging('user-repository')
 
 app = FastAPI()
 users_repository = UsersRepository()
@@ -17,13 +24,14 @@ def root():
     return {"message": "users service"}
 
 
-@app.get("/users/{user_id}")
-def get_users(user_id: int):
+@app.get("/users/me")
+def get_users(user_id: Annotated[int, Depends(get_current_user_id)]):
     return users_controller.handle_get_user(user_id)
 
 
 @app.get("/users")
-def get_all_users(ids: Annotated[Union[List[str], None], Query()] = None):
+def get_all_users(_: Annotated[int, Depends(get_current_user_id)],
+                  ids: Annotated[Union[List[str], None], Query()] = None):
     return users_controller.handle_get_all_users(ids)
 
 
@@ -39,6 +47,16 @@ async def login_with_google(request: LoginRequest):
 
 @app.patch("/users/me")
 async def update_user(update_data: UpdateUserSchema,
-                      request: Request):
-    return users_controller.handle_update_user(update_data.dict(),
-                                               request)
+                      user_id: Annotated[int, Depends(get_current_user_id)]):
+    return users_controller.handle_update_user(update_data.dict(), user_id)
+
+
+@app.post("/users/me/notification")
+async def create_notification(
+    user_id: Annotated[int,
+                       Depends(get_current_user_id)],
+    create_notification: CreateNotificationSchema
+):
+    logger.info(f"Creating notification for user {user_id}")
+    return users_controller.handle_create_notification(
+        create_notification.dict(), user_id)

app/models/alarms.py

@@ -16,13 +16,13 @@ class Alarm(Base):
                                     autoincrement=True)
     id_user: Mapped[int] = mapped_column(ForeignKey(f"{SCHEMA}.users.id"),
                                          nullable=False, unique=True)
-    datetime: Mapped[TIMESTAMP] = mapped_column(TIMESTAMP, nullable=False)
+    date_time: Mapped[TIMESTAMP] = mapped_column(TIMESTAMP, nullable=False)
     content: Mapped[str] = mapped_column(String(128), nullable=False)
 
     def __repr__(self) -> str:
         return (
             f"Alarm(id={self.id}, "
             f"id_user={self.id_user}, "
-            f"datetime={self.datetime}, "
+            f"date_time={self.date_time}, "
             f"content={self.content})"
         )

app/repository/Users.py

@@ -96,10 +96,12 @@ def edit_user(self, user_id: int, data_to_edit: dict):
         self.session.commit()
         return user
 
-    def __parse_result(self, result):
-        if not result:
-            return []
-        return [r.__dict__ for r in result]
+    @withSQLExceptionsHandle()
+    def create_notification(self, user_id: int, notification_data: dict):
+        alarm = Alarm(**notification_data, id_user=user_id)
+        self.session.add(alarm)
+        self.session.commit()
+        return alarm
 
     @withSQLExceptionsHandle()
     def get_users_to_notify(self,
@@ -128,3 +130,8 @@ def get_users_to_notify(self,
                             .filter(Alarm.datetime == date_time)
         result = query.all()
         return result
+
+    def __parse_result(self, result):
+        if not result:
+            return []
+        return [r.__dict__ for r in result]

app/schemas/Schemas.py

@@ -1,6 +1,7 @@
 from datetime import date
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, validator
 from typing import Optional, Dict
+from datetime import datetime
 
 
 class UserSchema(BaseModel):
@@ -40,3 +41,14 @@ class UpdateUserSchema(BaseModel):
     nickname: Optional[str] = None
     biography: Optional[str] = None
     device_token: Optional[str] = None
+
+
+class CreateNotificationSchema(BaseModel):
+    date_time: datetime = Field(..., alias='date_time')
+    content: str = Field(..., max_length=128)
+
+    @validator('date_time')
+    def validate_date_time(cls, v):
+        if v.minute % 5 != 0:
+            raise ValueError('Minutes must be multiples of 5')
+        return v

app/security/JWTBearer.py

@@ -0,0 +1,57 @@
+from typing import Annotated
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi import Request, Depends, HTTPException, status
+from jose import JWTError, jwt
+from pydantic import BaseModel
+from os import environ
+
+JWT_SECRET = environ.get("JWT_SECRET")
+HASH_ALGORITHM = environ.get("HASH_ALGORITHM")
+TOKEN_FIELD_NAME = "x-access-token"
+
+
+class JWTBearer(HTTPBearer):
+    def __init__(self, auto_error: bool = True):
+        super().__init__(auto_error=auto_error)
+
+    async def __call__(self, request: Request):
+        try:
+            credentials: HTTPAuthorizationCredentials = await super().\
+                __call__(request)
+            if credentials.scheme != "Bearer":
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="(Bearer) Invalid scheme or token.",
+                )
+            return credentials.credentials
+
+        except Exception:
+            access_token = request.headers.get(TOKEN_FIELD_NAME)
+            if access_token:
+                return access_token
+
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="(x-access-token) Invalid scheme or token.",
+            )
+
+
+class TokenData(BaseModel):
+    user_id: int | None = None
+
+
+async def get_current_user_id(token: Annotated[dict, Depends(JWTBearer())]):
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, JWT_SECRET, algorithms=[HASH_ALGORITHM])
+        user_id: int = payload.get("user_id")
+        if user_id is None:
+            raise credentials_exception
+        token_data = TokenData(user_id=user_id)
+    except JWTError:
+        raise credentials_exception
+    return token_data.user_id

app/service/Users.py

@@ -52,6 +52,14 @@ def update_user(self, user_id: int, update_data: dict):
             self.user_repository.rollback()
             raise e
 
+    def create_notification(self, user_id: int, notification_data: dict):
+        try:
+            self.user_repository.create_notification(user_id,
+                                                     notification_data)
+        except Exception as e:
+            self.user_repository.rollback()
+            raise e
+
     def _generate_nickname(self, name):
 
         name_without_spaces = name.replace(" ", "")

requirements.txt

@@ -9,4 +9,5 @@ pyjwt
 httpx
 APScheduler>=3.10,<4.0
 arq
-pytz
+pytz
+python-jose