Refactor database connection handling to use context manager (#22)

* Refactor database connection handling to use context manager for cursor operations

* Enhance database cursor management with rollback on exception and remove redundant commits in authentication and picture routes

* Refactor error handling in registration to improve readability and maintainability

* Simplify error response in login function to enhance security and user experience

* Remove unused database connection import from multiple route files

Author: Vianpyro

db.py

@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 import pymysql.cursors
 from flask import current_app
 
@@ -10,3 +12,16 @@ def get_db_connection():
         database=current_app.config["MYSQL_DB"],
         cursorclass=pymysql.cursors.DictCursor,
     )
+
+
+@contextmanager
+def database_cursor():
+    db = get_db_connection()
+    try:
+        yield db.cursor()
+        db.commit()
+    except Exception as e:
+        db.rollback()
+        raise e
+    finally:
+        db.close()

routes/authentication.py

@@ -6,7 +6,7 @@
 from flask import Blueprint, jsonify, request
 from pymysql import MySQLError
 
-from db import get_db_connection
+from db import database_cursor
 from jwt_helper import (
     TokenError,
     extract_token_from_header,
@@ -42,6 +42,25 @@ def validate_password(password):
     )
 
 
+def get_person_by_email(email):
+    with database_cursor() as cursor:
+        cursor.callproc("login_person", (email,))
+        return cursor.fetchone()
+
+
+def verify_password(password, stored_password, salt):
+    seasoned_password = password.encode("utf-8") + salt + PEPPER
+    try:
+        return ph.verify(stored_password, seasoned_password)
+    except exceptions.VerifyMismatchError:
+        return False
+
+
+def update_last_login(person_id):
+    with database_cursor() as cursor:
+        cursor.callproc("update_last_login", (person_id,))
+
+
 @authentication_blueprint.route("/register", methods=["POST"])
 def register():
     data = request.get_json()
@@ -58,23 +77,19 @@ def register():
 
     hashed_password, salt = hash_password_with_salt_and_pepper(password)
 
-    db = get_db_connection()
-    with db.cursor() as cursor:
-        try:
+    try:
+        with database_cursor() as cursor:
             cursor.callproc(
                 "register_person", (name, email, hashed_password, salt, language_code)
             )
-            db.commit()
-        except MySQLError as e:
-            # Check for specific error messages in the SQL error
-            if "User name already exists" in str(e):
-                return jsonify(message="User name already exists"), 400
-            elif "Email already exists" in str(e):
-                return jsonify(message="Email already exists"), 400
-            else:
-                return jsonify(message="An error occurred during registration"), 500
-
-    db.close()
+    except MySQLError as e:
+        if "User name already exists" in str(e):
+            return jsonify(message="User name already exists"), 400
+        elif "Email already exists" in str(e):
+            return jsonify(message="Email already exists"), 400
+        else:
+            return jsonify(message="An error occurred during registration"), 500
+
     return jsonify(message="User created successfully"), 201
 
 
@@ -87,30 +102,26 @@ def login():
     if not email or not password:
         return jsonify(message="Email and password are required"), 400
 
-    db = get_db_connection()
-    with db.cursor() as cursor:
-        cursor.callproc("login_person", (email,))
-        person = cursor.fetchone()
-
-        if not person:
-            return jsonify(message="Invalid credentials"), 401
+    person = get_person_by_email(email)
 
-        person_id = person["person_id"]
-        stored_password = person["hashed_password"]
-        salt = person["salt"]
-        seasoned_password = password.encode("utf-8") + salt + PEPPER
-
-        try:
-            ph.verify(stored_password, seasoned_password)
-            access_token = generate_access_token(person_id)
-            refresh_token = generate_refresh_token(person_id)
-            return jsonify(
-                message="Login successful",
-                access_token=access_token,
-                refresh_token=refresh_token,
-            )
-        except exceptions.VerifyMismatchError:
+    try:
+        if not person or not verify_password(
+            password, person["hashed_password"], person["salt"]
+        ):
             return jsonify(message="Invalid credentials"), 401
+    except Exception:
+        return jsonify(message="An internal error occurred"), 500
+
+    person_id = person["person_id"]
+    access_token = generate_access_token(person_id)
+    refresh_token = generate_refresh_token(person_id)
+    update_last_login(person_id)
+
+    return jsonify(
+        message="Login successful",
+        access_token=access_token,
+        refresh_token=refresh_token,
+    )
 
 
 @authentication_blueprint.route("/refresh", methods=["POST"])

routes/comment.py

@@ -1,55 +1,45 @@
 from flask import Blueprint, jsonify
 
-from db import get_db_connection
+from db import database_cursor
 
 comment_blueprint = Blueprint("comment", __name__)
 
 
 @comment_blueprint.route("/all", methods=["GET"])
 def get_all_comments():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_comments")
         comments = cursor.fetchall()
-    db.close()
     return jsonify(comments)
 
 
 @comment_blueprint.route("/<int:comment_id>", methods=["GET"])
 def get_comment_by_id(comment_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_comment_by_id", (comment_id,))
         comment = cursor.fetchone()
-    db.close()
     return jsonify(comment)
 
 
 @comment_blueprint.route("/person/<int:person_id>", methods=["GET"])
 def get_comments_by_person(person_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_comments_by_person", (person_id,))
         comments = cursor.fetchall()
-    db.close()
     return jsonify(comments)
 
 
 @comment_blueprint.route("/recipe/<int:recipe_id>", methods=["GET"])
 def get_comments_by_recipe(recipe_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_comments_by_recipe", (recipe_id,))
         comments = cursor.fetchall()
-    db.close()
     return jsonify(comments)
 
 
 @comment_blueprint.route("/count/recipe/<int:recipe_id>", methods=["GET"])
 def get_comment_count_by_recipe(recipe_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_comment_count_by_recipe", (recipe_id,))
         count = cursor.fetchone()
-    db.close()
     return jsonify(count)

routes/ingredient.py

@@ -1,25 +1,21 @@
 from flask import Blueprint, jsonify
 
-from db import get_db_connection
+from db import database_cursor
 
 ingredient_blueprint = Blueprint("ingredient", __name__)
 
 
 @ingredient_blueprint.route("/all", methods=["GET"])
 def get_all_ingredients():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_ingredients")
         ingredients = cursor.fetchall()
-    db.close()
     return jsonify(ingredients)
 
 
 @ingredient_blueprint.route("/<int:ingredient_id>", methods=["GET"])
 def get_ingredient_by_id(ingredient_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_ingredient_by_id", (ingredient_id,))
         ingredient = cursor.fetchone()
-    db.close()
     return jsonify(ingredient)

routes/language.py

@@ -1,27 +1,23 @@
 from flask import Blueprint, jsonify
 
-from db import get_db_connection
+from db import database_cursor
 
 language_blueprint = Blueprint("language", __name__)
 
 
 @language_blueprint.route("/all", methods=["GET"])
 def get_all_languages():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_languages")
         languages = cursor.fetchall()
-    db.close()
     return jsonify(languages)
 
 
 @language_blueprint.route("/<int:language_id>", methods=["GET"])
 def get_language_by_id(language_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_language_by_id", (language_id,))
         language = cursor.fetchone()
-    db.close()
     return jsonify(language)
 
 
@@ -30,29 +26,23 @@ def get_language_by_language_code(language_code):
     if len(language_code) != 2:
         return jsonify({"error": "Invalid language code"}), 400
 
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_language_by_code", (language_code,))
         language = cursor.fetchone()
-    db.close()
     return jsonify(language)
 
 
 @language_blueprint.route("/used", methods=["GET"])
 def get_used_languages():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_languages_with_users")
         languages = cursor.fetchall()
-    db.close()
     return jsonify(languages)
 
 
 @language_blueprint.route("/stats", methods=["GET"])
 def get_language_stats():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_languages_usage_statistics")
         stats = cursor.fetchall()
-    db.close()
     return jsonify(stats)

routes/person.py

@@ -1,25 +1,21 @@
 from flask import Blueprint, jsonify
 
-from db import get_db_connection
+from db import database_cursor
 
 person_blueprint = Blueprint("person", __name__)
 
 
 @person_blueprint.route("/all", methods=["GET"])
 def get_all_persons():
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_all_persons")
         persons = cursor.fetchall()
-    db.close()
     return jsonify(persons)
 
 
 @person_blueprint.route("/<int:person_id>", methods=["GET"])
 def get_person_by_id(person_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_person_by_id", (person_id,))
         person = cursor.fetchone()
-    db.close()
     return jsonify(person)

routes/picture.py

@@ -3,7 +3,7 @@
 from flask import Blueprint, jsonify, request, send_from_directory
 
 from config import PICTURE_FOLDER
-from db import get_db_connection
+from db import database_cursor
 from jwt_helper import token_required
 
 picture_blueprint = Blueprint("picture", __name__)
@@ -23,34 +23,28 @@ def extract_file_extension(filename):
 def get_all_pictures():
     picture_type = request.args.get("type", None)
 
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc(
             "get_all_pictures" if picture_type is None else "get_pictures_by_type",
             (picture_type,),
         )
         pictures = cursor.fetchall()
-    db.close()
     return jsonify(pictures)
 
 
 @picture_blueprint.route("/<int:picture_id>", methods=["GET"])
 def get_picture_by_id(picture_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_picture_by_id", (picture_id,))
         picture = cursor.fetchone()
-    db.close()
     return jsonify(picture)
 
 
 @picture_blueprint.route("/author/<int:author_id>", methods=["GET"])
 def get_pictures_by_author(author_id):
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc("get_pictures_by_author", (author_id,))
         picture = cursor.fetchall()
-    db.close()
     return jsonify(picture)
 
 
@@ -83,11 +77,8 @@ def upload_picture():
         case _:
             return jsonify({"error": f"Invalid picture type: {picture_type}"}), 400
 
-    db = get_db_connection()
-    with db.cursor() as cursor:
+    with database_cursor() as cursor:
         cursor.callproc(procedure, (hexname, request.person_id))
-        db.commit()
-    db.close()
 
     fullpath = os.path.normpath(os.path.join(PICTURE_FOLDER, hexname))
     if not fullpath.startswith(PICTURE_FOLDER):