Limit the number of screenshots (#2)

Adding parameter to limit number of screenshots --------- Co-authored-by: Sebastian Gomez-Gonzalez <[email protected]>
sebasutp · Oct 4, 2024 · d063e3b · d063e3b
1 parent 5b469b8
commit d063e3b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/app/api.py b/app/api.py
@@ -1,6 +1,6 @@
 """API methods."""
 
-from datetime import timedelta
+from datetime import timedelta, datetime
 from typing import Annotated
 
 from fastapi import Body, Depends, FastAPI, HTTPException
@@ -68,7 +68,7 @@ async def login(
 
 @app_obj.get("/screenshots", tags=["screenshots"], response_model=list[model.Screenshot])
 async def get_screenshots(
-    *,
+    num_screenshots: int = 3,
     session: Session = Depends(model.get_db_session),
     current_user_id: model.UserId = Depends(auth_handler.get_current_user_id)):
     """
@@ -90,7 +90,9 @@ async def get_screenshots(
         A list of `model.Screenshot` objects representing the user's screenshots.
     """
     statement = select(model.Screenshot).where(
-        model.Screenshot.owner_id == current_user_id.id)
+        model.Screenshot.owner_id == current_user_id.id).order_by(
+            model.Screenshot.created_on.desc()
+        ).limit(num_screenshots)
     return session.exec(statement)
 
 
@@ -154,7 +156,10 @@ async def add_screenshots(
         A `model.Screenshot` object representing the newly created screenshot.
     """
     screenshot_db = model.Screenshot.model_validate(
-        screenshot, update={"owner_id": current_user_id.id})
+        screenshot, update={
+            "owner_id": current_user_id.id,
+            "created_on": datetime.now()
+            })
     screenshot_db.external_id = crypto.generate_random_base64_string(32)
     session.add(screenshot_db)
     session.commit()

diff --git a/app/model.py b/app/model.py
@@ -2,6 +2,7 @@
 """
 
 import os
+from datetime import datetime
 from typing import Optional, Union
 
 from pydantic import BaseModel, EmailStr
@@ -59,6 +60,7 @@ class ScreenshotCreate(SQLModel):
 class ScreenshotBase(ScreenshotCreate):
     owner_id: int | None = Field(default=None, foreign_key="user.id")
     external_id: str = Field(default = None)
+    created_on: datetime
     #owner: User | None = Relationship(back_populates="screenshots")
 
 class Screenshot(ScreenshotBase, table=True):