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)