Birthdays

functions/database/birthdays.py

@@ -1,27 +1,26 @@
-from functions.database import utils
+from database.db import session
+from database.models import Birthday
+from typing import Optional, List
 
 
-def get_user(userid):
-    connection = utils.connect()
-    cursor = connection.cursor()
-    cursor.execute("SELECT day, month, year FROM birthdays WHERE userid = %s", (int(userid),))
-    return cursor.fetchall()
+def get_user(userid: int) -> Optional[Birthday]:
+    return session.query(Birthday).filter(Birthday.userid == userid).scalar()
 
 
-def get_users_on_date(day, month):
-    connection = utils.connect()
-    cursor = connection.cursor()
-    cursor.execute("SELECT userid FROM birthdays WHERE day = %s AND month = %s", (int(day), int(month),))
-    return cursor.fetchall()
+def get_users_on_date(day: int, month: int) -> List[Birthday]:
+    return session.query(Birthday.userid).filter(Birthday.day == day and Birthday.month == month).all()
 
 
-def add_user(userid, day, month, year):
-    connection = utils.connect()
-    cursor = connection.cursor()
-    if get_user(userid):
-        cursor.execute("UPDATE birthdays SET day = %s, month = %s, year = %s WHERE userid = %s",
-                       (int(day), int(month), int(year), int(userid),))
+def add_user(userid: int, day: int, month: int, year: int):
+    bd: Optional[Birthday] = get_user(userid)
+
+    # Update user if they exist, otherwise insert entry
+    if bd is not None:
+        bd.day = day
+        bd.month = month
+        bd.year = year
     else:
-        cursor.execute("INSERT INTO birthdays(userid, day, month, year) VALUES (%s, %s, %s, %s)",
-                       (int(userid), int(day), int(month), int(year),))
-    connection.commit()
+        entry = Birthday(userid=userid, day=day, month=month, year=year)
+        session.add(entry)
+
+    session.commit()