diff --git a/api/utils.py b/api/utils.py index b44b742..beffffb 100644 --- a/api/utils.py +++ b/api/utils.py @@ -19,6 +19,14 @@ def validate_access(owner_id, message='Access denied.'): abort(403, message) +def get_user_or_404(user_id): + "Get user from database or abort 404" + user = db.session.get(User, user_id) + if user is None: + abort(404, "User not found") + return user + + def init_db(): """Create default admin account if database is empty""" with db.session.begin(): diff --git a/api/views.py b/api/views.py index ecf68ba..4d73b69 100644 --- a/api/views.py +++ b/api/views.py @@ -2,7 +2,7 @@ from flask import Blueprint, jsonify, request, abort from flask_jwt_extended import create_access_token, set_access_cookies, jwt_required, \ verify_jwt_in_request, get_jwt_identity, unset_jwt_cookies, get_jwt from models import db, RevokedToken, User -from utils import admin_required, validate_access +from utils import admin_required, validate_access, get_user_or_404 from werkzeug.security import check_password_hash, generate_password_hash user_bp = Blueprint('user_bp', __name__) @@ -23,9 +23,7 @@ def get_all_users(): @jwt_required() def get_user(user_id): validate_access(user_id) # check if user tries to read other user account details - user = db.session.get(User, user_id) - if user is None: - abort(404, "User not found.") + user = get_user_or_404(user_id) return jsonify(user.to_dict()) @@ -59,7 +57,7 @@ def edit_user(user_id): if request_fields != editable_fields: abort(400, "Invalid request data structure.") - user_to_update = User.query.get_or_404(user_id) + user_to_update = get_user_or_404(user_id) for field_name in editable_fields: requested_value = request_data.get(field_name) if requested_value is None: @@ -75,9 +73,7 @@ def edit_user(user_id): @jwt_required() def remove_user(user_id): validate_access(user_id) # Only admin can remove other users accounts - user_to_delete = db.session.get(User, user_id) - if user_to_delete is None: - abort(404, "User not found.") + user_to_delete = get_user_or_404(user_id) db.session.delete(user_to_delete) db.session.commit() return jsonify({"msg": "User removed successfully."})