From 32ac22806c14009ae7a01104d4c2497c3e44faf6 Mon Sep 17 00:00:00 2001
From: Derek Wang <derek.kiah.wang@gmail.com>
Date: Mon, 3 Mar 2025 19:28:15 -0800
Subject: [PATCH] Fixed POST requests in system

---
 app/database.py          | 19 +++++++++++++++----
 app/main.py              |  8 ++++++--
 app/models/models.py     |  9 +++++++++
 tests/varied_workload.py | 15 ++++++++++-----
 4 files changed, 40 insertions(+), 11 deletions(-)
 create mode 100644 app/models/models.py

diff --git a/app/database.py b/app/database.py
index aa55595..747e1b7 100644
--- a/app/database.py
+++ b/app/database.py
@@ -1,12 +1,16 @@
-
 from tinydb import TinyDB, Query
 from config import DB_FILE
+import shutil
+
 import random
 
 DB_LOCATION = "database/datastore/" + DB_FILE
 
 # Initialize TinyDB as a NoSQL key-value store
-db = TinyDB(DB_LOCATION)
+# We don't want to change our main DB file, so we will make a temp DB file and use that as our DB file
+shutil.copyfile( DB_LOCATION, "temp_DB.json" )
+
+db = TinyDB("temp_DB.json")
 User = Query()
 
 def get_user_ids():
@@ -33,9 +37,16 @@ def get_user_profile(user_id):
     result = db.search(User.user_id == user_id)
     return result[0] if result else None
 
-def update_user_profile(user_id, name, followers, bio, posts, friends):
+def update_user_profile( data ):
     """Update user profile in TinyDB"""
-    db.upsert({"user_id": user_id, "name": name, "followers": followers, "bio": bio, "posts": posts, "friends": friends}, User.user_id == user_id)
+    user_id = str( data["user_id"] )
+
+    # Basically make sure friends stay the same (for prefetching). Not great implementation, but it works
+    curr_user = db.search(User.user_id == user_id)
+    if( curr_user and data["friends"] == None ):
+        data["friends"] = curr_user[0]["friends"]
+
+    db.upsert( data, User.user_id == user_id )
 
 def init_db():
     """Ensure TinyDB is initialized before FastAPI starts and prepopulate some data"""
diff --git a/app/main.py b/app/main.py
index a929695..28440a6 100644
--- a/app/main.py
+++ b/app/main.py
@@ -8,6 +8,7 @@ from cache.nocache import NoCache
 from cache.idealcache import IdealCache
 from cache.read_after_write_cache import ReadAfterWriteCache
 from config import CACHE_STRATEGY, CACHE_LIMIT, L2_CACHE_LIMIT
+from models.models import User
 import time
 
 app = FastAPI()
@@ -62,8 +63,11 @@ def fetch_user_profile(user_id: str):
     return {"user_id": user_id, "profile": profile, "source": "database", "time_ms": (time.time() - start) * 1000}
 
 @app.post("/update_user/")
-def modify_user_profile(user_id: str, name: str, followers: int, bio: str, posts: str, friends: list[str]):
+async def modify_user_profile(user_data : User):
     """Update user profile and refresh cache"""
-    update_user_profile(user_id, name, followers, bio, posts, friends)
+    user_id=user_data.user_id
+    user_dict = user_data.dict()
+    
+    update_user_profile(user_dict)
     cache.invalidate(user_id)  # Invalidate old cache
     return {"message": "User profile updated successfully"}
diff --git a/app/models/models.py b/app/models/models.py
new file mode 100644
index 0000000..a424ba7
--- /dev/null
+++ b/app/models/models.py
@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+class User(BaseModel):
+    user_id: str
+    name: str | None = None
+    followers: int | None = None
+    bio: str | None = None
+    posts: str | None = None
+    friends: list | None = None
\ No newline at end of file
diff --git a/tests/varied_workload.py b/tests/varied_workload.py
index 63eed3c..b3d1928 100644
--- a/tests/varied_workload.py
+++ b/tests/varied_workload.py
@@ -30,12 +30,13 @@ def generate_request(workload, last_updated=None):
     """Generate read or write requests based on workload type"""
     if random.random() < workload["read"]:
         user_id = select_user(workload, last_updated)
-        return baseurl + f"/user/{user_id}", "GET"
+        return baseurl + f"/user/{user_id}", None, "GET"
 
     # Write operation (updates user profile)
     user_id = select_user(workload, last_updated)
-    url = baseurl + f"/update_user/?user_id={user_id}&name=UpdatedUser&followers=500&bio=Updated&posts=UpdatedPost"
-    return url, "POST"
+    write_obj = { "user_id":user_id,"name": "UpdatedUser", "followers":"500","bio":"Updated","posts":"UpdatedPost"}
+    url = baseurl + f"/update_user/"
+    return url, write_obj, "POST"
 
 def select_user(workload, last_updated):
     """Selects a user based on workload type"""
@@ -55,8 +56,12 @@ def run_workload(name, workload):
     last_updated = None
 
     for _ in tqdm(range(10000), desc=f"Running {name}"):
-        url, method = generate_request(workload, last_updated)
-        response = requests.request(method, url)
+        url, data, method = generate_request(workload, last_updated)
+        
+        if( method == "GET" ):
+            response = requests.request(method, url)
+        else:
+            response = requests.post(url, json = data)
 
         try:
             content = json.loads(response.content)