From 55e7ee145c470e832b56c9540f520aedf8792c15 Mon Sep 17 00:00:00 2001
From: Parker <contact@pkrm.dev>
Date: Thu, 19 Dec 2024 16:35:18 -0600
Subject: [PATCH] Add database cleanup scheduler

---
 api/main.py          | 22 ++++++++++++++++++++++
 api/util/clean_db.py | 25 +++++++++++++++++++++++++
 requirements.txt     |  3 ++-
 3 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 api/util/clean_db.py

diff --git a/api/main.py b/api/main.py
index 0e45e8a..5e38a1c 100644
--- a/api/main.py
+++ b/api/main.py
@@ -6,10 +6,31 @@ from api.routes.links_routes import router as links_router
 from api.routes.user_routes import router as user_router
 from api.routes.log_routes import router as log_router
 from typing import Annotated
+from apscheduler.schedulers.background import BackgroundScheduler
+from contextlib import asynccontextmanager
 
 from api.util.db_dependency import get_db
+from api.util.clean_db import clean_db
 from api.util.log import log
 from models import Link
+from config import LOG
+
+scheduler = BackgroundScheduler()
+
+
+# Handle the database cleanup scheduler
+def start_scheduler():
+    scheduler.add_job(clean_db, "interval", days=1)
+    scheduler.start()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    try:
+        start_scheduler()
+        yield
+    finally:
+        scheduler.shutdown()
 
 
 app = FastAPI(
@@ -21,6 +42,7 @@ app = FastAPI(
         "identifier": "Unlicense",
         "url": "https://unlicense.org",
     },
+    lifespan=lifespan,
 )
 
 app.add_middleware(
diff --git a/api/util/clean_db.py b/api/util/clean_db.py
new file mode 100644
index 0000000..11438aa
--- /dev/null
+++ b/api/util/clean_db.py
@@ -0,0 +1,25 @@
+import datetime
+
+from api.util.db_dependency import get_db
+from models import Link, Log
+
+"""
+Remove expired short links and their associated logs
+"""
+
+
+def clean_db():
+    db = next(get_db())
+    # Get all expired short links
+    expired_links = (
+        db.query(Link)
+        .filter(Link.expire_date < datetime.datetime.today())
+        .all()
+    )
+
+    # Delete all expired short links and their logs
+    for link in expired_links:
+        logs = db.query(Log).filter(Log.link == link.link).all()
+        for log in logs:
+            db.delete(log)
+        db.delete(link)
diff --git a/requirements.txt b/requirements.txt
index 0101270..0b5a157 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,4 +9,5 @@ bcrypt==4.1.3
 SQLAlchemy==2.0.31
 validators==0.28.3
 requests==2.32.3
-ua-parser==0.18.0
\ No newline at end of file
+ua-parser==0.18.0
+APScheduler==3.11.0
\ No newline at end of file