Merge pull request #15 from Police-Data-Accessibility-Project/mc_repl…

…ace_with_abort

Move from flask_restful to flask_restx. Replace error dictionaries with abort

app.py

@@ -1,5 +1,5 @@
 from flask import Flask
-from flask_restful import Api
+from flask_restx import Api
 from flask_cors import CORS
 from resources.User import User
 from resources.Login import Login

middleware/security.py

@@ -1,4 +1,8 @@
 import functools
+from flask import request, jsonify
+from flask_restx import abort
+
+
 from collections import namedtuple
 
 from http import HTTPStatus

requirements.txt

@@ -20,6 +20,7 @@ exceptiongroup==1.1.3
 Flask==2.3.2
 Flask-Cors==4.0.0
 Flask-RESTful==0.3.10
+flask-restx==1.3.0
 Flask-SQLAlchemy~=3.1.1
 gotrue==1.0.3
 gunicorn==21.2.0

resources/Login.py

@@ -1,3 +1,4 @@
+from flask_restx import abort
 from werkzeug.security import check_password_hash
 from flask import request
 from middleware.login_queries import login_results, create_session_token
@@ -35,4 +36,4 @@ def post(self):
                 "data": token,
             }
 
-        return {"message": "Invalid email or password"}, 401
+        abort(code=401, message="Invalid email or password")

resources/PsycopgResource.py

@@ -1,8 +1,7 @@
 import functools
 from typing import Callable, Any, Union, Tuple, Dict
 
-from flask import make_response
-from flask_restful import Resource
+from flask_restx import abort, Resource
 
 
 def handle_exceptions(
@@ -37,15 +36,16 @@ def wrapper(
         except Exception as e:
             self.psycopg2_connection.rollback()
             print(str(e))
-            return make_response({"message": str(e)}, 500)
+            abort(http_status_code=500, message=str(e))
 
     return wrapper
 
 
 class PsycopgResource(Resource):
-    def __init__(self, **kwargs):
+    def __init__(self, *args, **kwargs):
         """
         Initializes the resource with a database connection.
         - kwargs (dict): Keyword arguments containing 'psycopg2_connection' for database connection.
         """
+        super().__init__(*args, **kwargs)
         self.psycopg2_connection = kwargs["psycopg2_connection"]

resources/RefreshSession.py

@@ -1,5 +1,7 @@
 from flask import request
 
+from flask_restx import abort
+
 from middleware.custom_exceptions import TokenNotFoundError
 from middleware.login_queries import (
     get_session_token_user_data,
@@ -40,3 +42,5 @@ def post(self) -> Dict[str, Any]:
             "message": "Successfully refreshed session token",
             "data": token,
         }
+
+        abort(code=403, message="Invalid session token")

resources/ResetPassword.py

@@ -1,3 +1,4 @@
+from flask_restx import abort
 from werkzeug.security import generate_password_hash
 from flask import request
 from middleware.reset_token_queries import (
@@ -32,13 +33,13 @@ def post(self) -> Dict[str, Any]:
         token_data = check_reset_token(cursor, token)
         email = token_data.get("email")
         if "create_date" not in token_data:
-            return {"message": "The submitted token is invalid"}, 400
+            abort(code=400, message="The submitted token is invalid")
 
         token_create_date = token_data["create_date"]
         token_expired = (dt.utcnow() - token_create_date).total_seconds() > 900
         delete_reset_token(cursor, token_data["email"], token)
         if token_expired:
-            return {"message": "The submitted token is invalid"}, 400
+            abort(code=400, message="The submitted token is invalid")
 
         password_digest = generate_password_hash(password)
         cursor = self.psycopg2_connection.cursor()

resources/ResetTokenValidation.py

@@ -1,4 +1,6 @@
 from flask import request
+from flask_restx import abort
+
 from middleware.reset_token_queries import (
     check_reset_token,
 )
@@ -22,6 +24,6 @@ def post(self):
         token_expired = (dt.utcnow() - token_create_date).total_seconds() > 900
 
         if token_expired:
-            return {"message": "The submitted token is invalid"}, 400
+            abort(code=400, message="The submitted token is invalid")
 
         return {"message": "Token is valid"}

tests/resources/test_RefreshSession.py

@@ -125,9 +125,7 @@ def test_post_refresh_session_unexpected_error(
     )
 
     check_response_status(response, 500)
-    assert response.json == {
-        "message": "An unexpected error occurred",
-    }
+    assert response.json["message"] == "An unexpected error occurred"
     mock_get_session_token_user_data.assert_called_once_with(
         mock_cursor, "old_test_session_token"
     )