Refactor database interactions to use SQLAlchemy ORM; update user and token management in api_security.py and user.py; add models for Items, Tokens, and User

auth/api_security.py

@@ -7,6 +7,7 @@ import logging
 from argon2 import PasswordHasher
 import random
 import string
+from models import Tokens
 
 
 #Logging
@@ -23,17 +24,7 @@ try:
 except Exception as e:
     logging.exception('An error occurred')
 
-load_dotenv()
+session = db.session
-db_name = os.getenv("DB_NAME")
-db_psswd = os.getenv("DB_PASSWORD")
-db_user = os.getenv("DB_USER")
-db_port = os.getenv("DB_PORT") or "3306"
-db_host = os.getenv("DB_HOST")
-if db_host == None:
-    logging.error("no hostname provided")
-    os._exit(1)
-
-db.connect(db_user, db_psswd, db_host, int(db_port), db_name)
 
 
 def generate_token(duration):
@@ -43,11 +34,9 @@ def generate_token(duration):
     hash = ph.hash(token)
     end_date = enddate(duration)
     time = datetime.now()
-    cursor = db.connection.cursor()
+    new_item = Tokens(tokens=hash, time=time, end_date=end_date)
-    executed = "INSERT INTO tokens (tokens, time, end_date) VALUES (?, ?, ?)"
+    session.add(new_item)
-    executed_data = (hash, time, end_date)
+    session.commit()
-    cursor.execute(executed, executed_data)
-    db.connection.commit()
     logging.info(f"Token till {end_date} created")
     return token
 
@@ -68,10 +57,8 @@ def enddate(duration):
 
 def check_token(token):
     ph = PasswordHasher()
-    cursor = db.connection.cursor()
     valid = False
-    cursor.execute("SELECT tokens FROM tokens")
+    hashes = [row[0] for row in session.query(Tokens.tokens).all()]
-    hashes = [row[0] for row in cursor.fetchall()]
     for hash in hashes:
         try:
             valid = ph.verify(hash, token)
@@ -84,13 +71,12 @@ def check_token(token):
 
 
 def check_validaty():
-    cursor = db.connection.cursor()
+    end_dates = [row[0] for row in session.query(Tokens.end_date).all()]
-    cursor.execute("SELECT end_date FROM tokens")
-    end_dates = [row[0] for row in cursor.fetchall()]
     for end_date in end_dates:
         end_date_dt = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S")
         if datetime.now() > end_date_dt:
-            cursor.execute("DELETE FROM tokens WHERE end_date = ?", (end_date,))
+            session.query(Tokens).filter_by(end_date=end_date).delete()
+            session.commit()
 
 
 def admin_token():
@@ -99,11 +85,9 @@ def admin_token():
     hash = ph.hash(token)
     time = datetime.now()
     end_date = enddate("forever")
-    cursor = db.connection.cursor()
+    new_item = Tokens(tokens=hash, time=time, end_date=end_date)
-    executed = "INSERT INTO tokens (tokens, time, end_date) VALUES (?, ?, ?)"
+    session.add(new_item)
-    executed_data = (hash, time, end_date)
+    session.commit()
-    cursor.execute(executed, executed_data)
-    db.connection.commit()
     logging.info(f"Token till {end_date} created")
     return token
 
@@ -113,8 +97,7 @@ def check_admin_token(token):
     if len(token) == 64:
         ph = PasswordHasher()
         cursor = db.connection.cursor()        
-        cursor.execute("SELECT tokens FROM tokens")
+        hashes = [row[0] for row in session.query(Tokens.tokens).all()]
-        hashes = [row[0] for row in cursor.fetchall()]
         for hash in hashes:
             try:
                 valid = ph.verify(hash, token)
@@ -127,14 +110,12 @@ def check_admin_token(token):
 def delete_token(token):
     if not len(token) == 64:
         ph = PasswordHasher()
-        cursor = db.connection.cursor() 
+        hashes = [row[0] for row in session.query(Tokens.tokens).all()]
-        cursor.execute("SELECT tokens FROM tokens")
-        hashes = [row[0] for row in cursor.fetchall()]
         for hash in hashes:
             try:
                 valid = ph.verify(hash, token)
                 if valid == True:
-                    cursor.execute("DELETE FROM tokens WHERE tokens = ?", (hash,))
+                    session.query(Tokens).filter_by(tokens=hash).delete()
                     db.connection.commit()
                     logging.critical("token deleted")
             except Exception as e:

auth/user.py

@@ -1,9 +1,7 @@
 from argon2 import PasswordHasher
 import db
-from dotenv import load_dotenv
+from models import User
 import logging
-import os
-import mariadb
 #Logging
 
 logging.basicConfig(
@@ -18,31 +16,16 @@ try:
 except Exception as e:
     logging.exception('An error occurred')
 
-load_dotenv()
+session = db.session
-db_name = os.getenv("DB_NAME")
-db_psswd = os.getenv("DB_PASSWORD")
-db_user = os.getenv("DB_USER")
-db_port = os.getenv("DB_PORT") or "3306"
-db_host = os.getenv("DB_HOST")
-if db_host == None:
-    logging.error("no hostname provided")
-    os._exit(1)
 
-db.connect(db_user, db_psswd, db_host, int(db_port), db_name)
+def create_user(username, password, role):
-cursor = db.connection.cursor()
-
-def create_user(username, password):
     user = username
     passwd = PasswordHasher().hash(password)
-    cursor = db.connection.cursor()
+    new_item = User(username=user, password=passwd, role=role)
-    executed = "INSERT INTO user (username, password) VALUES (?, ?)"
-    executed_data = (user, passwd)
-    cursor.execute(executed, executed_data)
-    db.connection.commit()
 
 def check_login(username, password):
-    cursor = db.connection.cursor()
+    result = session.query(User.password).filter_by(username=username).first()
-    hash_passwd = cursor.execute("SELECT password FROM user WHERE username = ?", (username,))
+    hash_passwd = result[0] if result else None
     valid = False
     try:
         valid = PasswordHasher().verify(hash_passwd, password)
@@ -53,8 +36,8 @@ def check_login(username, password):
         return False
 
 def delete_user(username, password):
-    cursor = db.connection.cursor()
+    result = session.query(User.password).filter_by(username=username).first()
-    hash_password = cursor.execute("SELECT password FROM user WHERE username = ?", (username,))#
+    hash_password = result[0] if result else None
     valid = False
     try:
         valid = PasswordHasher().verify(hash_password, password)
@@ -62,19 +45,21 @@ def delete_user(username, password):
         logging.error(e)
     if valid == True:
         try:
-            cursor.execute("DELETE FROM user WHERE username = ?", (username,))
+            session.query(User).filter_by(username=username).delete()
-            db.connection.commit()
+            session.commit()
             logging.info(f"User {username} was successfully deleted")
         except Exception as e:
             logging.error(e)
 
 def check_username(username):
     try:
-        cursor.execute("SELECT 1 FROM user WHERE username = ?", (username,))
+        exists = session.query(User).filter_by(username=username).first() is not None
-        exists = cursor.fetchone() is not None
         return exists  # True if taken, False if available
     except db.IntegrityError as e:
         if e.errno == 1062:
             logging.error("Username taken")
             return True #If the username is not available it returns True
 
+def get_role(username):
+    result = session.query(User.role).filter_by(username=username).first()
+    return result[0] if result else None

db.py

@@ -1,7 +1,10 @@
-import mariadb
 import logging
 from dotenv import load_dotenv
 import os
+from models import Items
+from sqlalchemy import create_engine, func
+from sqlalchemy.orm import sessionmaker
+from urllib.parse import quote_plus
 
 #Logging
 
@@ -17,102 +20,6 @@ try:
 except Exception as e:
     logging.exception('An error occurred')
 
-
-def connect(db_user, db_psswd, db_host, db_port, db_name):
-    try:
-        global connection
-        connection = mariadb.connect(user=db_user,
-            password=db_psswd,
-            host=db_host,
-            port=db_port,
-            database=db_name)
-        print("Connected to MariaDB successfully!")
-        logging.info("Connected to MariaDB successfully!")
-        
-    except mariadb.Error as e:
-        print(f"Error connecting to DB: {e}")
-        logging.error("Error connecting")
-        logging.debug(f"Error connecting to DB: {e}")
-
-def new_part(cloth_type, weight, pieces):
-    if type(pieces) is not int:
-        print("only ints allowed")
-        logging.error("only ints allowed")
-    else:
-        cursor = connection.cursor()
-        executed = "INSERT INTO luggage (clothes, weight, pieces) VALUES (?, ?, ?)"
-        executed_data = (cloth_type, weight, pieces)
-        cursor.execute(executed, executed_data)
-        connection.commit()
-        logging.info("new Insert is in DB")
-
-def fetch_clothes():
-    cursor = connection.cursor()
-    cursor.execute("SELECT clothes FROM luggage")
-    clothes = [str(row[0]) for row in cursor.fetchall()]
-    logging.info("fetched clothes")
-    return clothes
-
-def update_weight(cloth_type, weight):
-    cursor = connection.cursor()
-    executed = "UPDATE luggage SET weight = ? WHERE clothes = ?"
-    executed_data = (weight, cloth_type)
-    cursor.execute(executed, executed_data)
-    connection.commit()
-    logging.info("weight updated")
-
-def update_pieces(cloth_type, pieces):
-    if type(pieces) is not int:
-        logging.error("only ints allowed")
-    else: 
-        cursor = connection.cursor()
-        executed = "UPDATE luggage SET pieces = ? WHERE clothes = ?"
-        executed_data = (pieces, cloth_type)
-        cursor.execute(executed, executed_data)
-        connection.commit()
-        logging.info("pieces updated")
-
-def fetch_weight(cloth_type):
-    cursor = connection.cursor()
-    cursor.execute("SELECT weight FROM luggage WHERE clothes = ?", (cloth_type,))
-    weight = cursor.fetchone()[0]
-    logging.info("fetched weight")
-    return weight
-
-def fetch_luggage_weight():
-    cursor = connection.cursor()
-    cursor.execute("SELECT sum(weight) FROM luggage")
-    weight = cursor.fetchone()[0]
-    logging.info("fetched luggage weight")
-    return weight
-
-def fetch_piece(cloth_type):
-    cursor = connection.cursor()
-    cursor.execute("SELECT pieces FROM luggage WHERE clothes = ?", (cloth_type,))
-    pieces = cursor.fetchone()[0]
-    logging.info("fetched luggage pieces")
-    return pieces
-
-def fetch_all_pieces():
-    cursor = connection.cursor()
-    cursor.execute("SELECT sum(pieces) FROM luggage")
-    pieces = cursor.fetchone()[0]
-    logging.info("fetched all luggage pieces")
-    return pieces
-
-def delete(cloth_type):
-    cursor = connection.cursor()
-    cursor.execute("DELETE FROM luggage WHERE clothes = ?;", (cloth_type,))
-    logging.critical(f"{cloth_type} was deleted")
-    connection.commit()
-
-def check_clothe(clothe_type):
-    cursor = connection.cursor()
-    cursor.execute("SELECT (SELECT clothes FROM luggage WHERE clothes = ? LIMIT 1) AS clothes", (clothe_type,))
-    check = cursor.fetchone()[0]
-    return check
-
-
 #using .env
 load_dotenv()
 db_name = os.getenv("DB_NAME")
@@ -123,4 +30,72 @@ db_host = os.getenv("DB_HOST")
 if db_host == None:
     logging.error("no hostname provided")
     os._exit(1)
-connect(db_user, db_psswd, db_host, int(db_port), db_name)
+
+engine = create_engine(f"mariadb+mariadbconnector://{db_user}:{quote_plus(db_psswd)}@{db_host}:{db_port}/{db_name}")
+Session = sessionmaker(bind=engine)
+session = Session()
+
+
+def new_part(cloth_type, weight, pieces):
+    existing = session.query(Items).filter_by(clothes=cloth_type).first()
+    if existing:
+        logging.debug("Item already inserted")
+    else:
+        if type(pieces) is not int:
+            print("only ints allowed")
+            logging.error("only ints allowed")
+        else:
+            new_item = Items(clothes=cloth_type, weight=weight, pieces=pieces)
+            session.add(new_item)
+            session.commit()
+            logging.info("new Insert is in DB")
+
+def update_weight(cloth_type, weight):
+    item = session.query(Items).filter_by(clothes=cloth_type).first()
+    item.weight = weight
+    session.commit()
+    logging.info("weight updated")
+
+def update_pieces(cloth_type, pieces):
+    if type(pieces) is not int:
+        logging.error("only ints allowed")
+    else: 
+        item = session.query(Items).filter_by(clothes=cloth_type).first()
+        item.pieces = pieces
+        session.commit()
+        logging.info("pieces updated")
+
+def delete(cloth_type):
+    logging.critical(f"{cloth_type} was deleted")
+    session.query(Items).filter_by(clothes=cloth_type).delete()
+    session.commit()
+
+
+def fetch_clothes():
+    logging.info("fetched clothes")
+    return [row[0] for row in session.query(Items.clothes).all()]
+
+def fetch_weight(cloth_type):
+    logging.info("fetched weight")
+    result = session.query(Items.weight).filter_by(clothes=cloth_type).first()
+    return result[0] if result else None
+
+def fetch_luggage_weight():
+    logging.info("fetched whole weight")
+    result = session.query(func.sum(Items.weight)).scalar()
+    return result if result else None
+
+def fetch_piece(cloth_type):
+    logging.info("fetched pieces")
+    result = session.query(Items.pieces).filter_by(clothes=cloth_type).first()
+    return result[0] if result else None
+
+def fetch_all_pieces():
+    logging.info("fetched all pieces")
+    result = session.query(Items.pieces).all()
+    return [row[0] for row in session.query(Items.pieces).all()]
+
+def check_clothe(clothe_type):
+    exists = session.query(Items.clothes).filter_by(clothes=clothe_type).first() is not None
+    return exists
+

models/__init__.py

@@ -0,0 +1 @@
+from .models import Items, Tokens, User, Base

models/models.py

@@ -0,0 +1,23 @@
+from sqlalchemy.orm import declarative_base
+from sqlalchemy import Column, String, Numeric, Integer
+
+
+Base = declarative_base()
+
+class Tokens(Base):
+    __tablename__ = "tokens"
+    tokens = Column(String(255), primary_key=True)
+    time = Column(String(255))
+    end_date = Column(String(255))
+
+class Items(Base):
+    __tablename__ = "luggage"
+    clothes = Column(String(50), primary_key=True)
+    weight = Column(Numeric(20, 2))
+    pieces = Column(Integer)
+
+class User(Base):
+    __tablename__ = "user"
+    username = Column(String(50), primary_key=True)
+    password = Column(String(255))
+    role = Column(String(50))

requirements.txt

@@ -2,3 +2,7 @@ mariadb == 1.1.14
 python-dotenv == 1.2.3
 fastapi[standard] == 0.141.1
 argon2-cffi == 25.1.0
+Jinja2==3.1.6
+MarkupSafe==3.0.3
+itsdangerous==2.2.0
+alembic==1.20.0