Update database.py

database.py

@@ -1,42 +1,81 @@
 import sqlite3
-from datetime import date
 import os
+from datetime import date
 
 DB_FILE = "users.db"
 
 
 def get_connection():
-    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
-    return conn
+    return sqlite3.connect(DB_FILE, check_same_thread=False)
+
+
+def recreate_database():
+    if os.path.exists(DB_FILE):
+        os.remove(DB_FILE)
 
 
 def create_tables():
-    conn = get_connection()
-    cur = conn.cursor()
 
-    cur.execute("""
-    CREATE TABLE IF NOT EXISTS users(
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        email TEXT UNIQUE,
-        password TEXT,
-        otp TEXT
-    )
-    """)
-
-    cur.execute("""
-    CREATE TABLE IF NOT EXISTS weight_history(
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        email TEXT,
-        date TEXT,
-        weight REAL
-    )
-    """)
+    try:
+        conn = get_connection()
+        cur = conn.cursor()
+
+        cur.execute("""
+        CREATE TABLE IF NOT EXISTS users(
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            email TEXT UNIQUE,
+            password TEXT,
+            otp TEXT
+        )
+        """)
+
+        cur.execute("""
+        CREATE TABLE IF NOT EXISTS weight_history(
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            email TEXT,
+            date TEXT,
+            weight REAL
+        )
+        """)
 
-    conn.commit()
-    conn.close()
+        conn.commit()
+        conn.close()
+
+    except sqlite3.DatabaseError:
+        # database corrupted → recreate
+        recreate_database()
+
+        conn = get_connection()
+        cur = conn.cursor()
 
+        cur.execute("""
+        CREATE TABLE users(
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            email TEXT UNIQUE,
+            password TEXT,
+            otp TEXT
+        )
+        """)
+
+        cur.execute("""
+        CREATE TABLE weight_history(
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            email TEXT,
+            date TEXT,
+            weight REAL
+        )
+        """)
+
+        conn.commit()
+        conn.close()
+
+
+# -----------------------
+# USER FUNCTIONS
+# -----------------------
 
 def register_user(email, password):
+
     conn = get_connection()
     cur = conn.cursor()
 
@@ -47,13 +86,16 @@ def register_user(email, password):
         )
         conn.commit()
         return True
+
     except:
         return False
+
     finally:
         conn.close()
 
 
 def login_user_db(email, password):
+
     conn = get_connection()
     cur = conn.cursor()
 
@@ -69,7 +111,12 @@ def login_user_db(email, password):
     return user
 
 
+# -----------------------
+# OTP
+# -----------------------
+
 def save_otp(email, otp):
+
     conn = get_connection()
     cur = conn.cursor()
 
@@ -83,6 +130,7 @@ def save_otp(email, otp):
 
 
 def verify_otp_db(email, otp):
+
     conn = get_connection()
     cur = conn.cursor()
 
@@ -98,7 +146,12 @@ def verify_otp_db(email, otp):
     return user
 
 
+# -----------------------
+# WEIGHT TRACKER
+# -----------------------
+
 def save_weight(email, weight):
+
     today = str(date.today())
 
     conn = get_connection()
@@ -114,6 +167,7 @@ def save_weight(email, weight):
 
 
 def get_weight_history(email):
+
     conn = get_connection()
     cur = conn.cursor()
 
@@ -122,8 +176,8 @@ def get_weight_history(email):
         (email,)
     )
 
-    history = cur.fetchall()
+    data = cur.fetchall()
 
     conn.close()
 
-    return history
+    return data