Updated functionalities

app/routes/facilities.py

@@ -1,4 +1,4 @@
-from flask import Blueprint, render_template, redirect, url_for, flash, request
+from flask import Blueprint, render_template, redirect, url_for, flash, request, abort
 from flask_login import login_required, current_user
 from app import db
 from app.models.facility import Facility, Area
@@ -53,7 +53,9 @@ def create_facility():
 @bp.route('/<int:facility_id>')
 @login_required
 def view_facility(facility_id):
-    facility = Facility.query.get_or_404(facility_id)
+    facility = db.session.get(Facility, facility_id)
+    if facility is None:
+        abort(404)
     if current_user.role == 'customer':
         cids = get_customer_scope(current_user) or []
         if facility_id not in cids:
@@ -66,7 +68,9 @@ def view_facility(facility_id):
 @login_required
 @supervisor_required
 def edit_facility(facility_id):
-    facility = Facility.query.get_or_404(facility_id)
+    facility = db.session.get(Facility, facility_id)
+    if facility is None:
+        abort(404)
     form = FacilityForm(obj=facility)
     projects = Project.query.filter_by(active=True).order_by(Project.name).all()
     form.project_id.choices = [(0, '— None —')] + [(p.id, p.name) for p in projects]
@@ -92,7 +96,9 @@ def edit_facility(facility_id):
 @login_required
 @admin_required
 def delete_facility(facility_id):
-    facility = Facility.query.get_or_404(facility_id)
+    facility = db.session.get(Facility, facility_id)
+    if facility is None:
+        abort(404)
     
     if facility.inspections.count() > 0:
         flash(f'Cannot delete "{facility.name}" — it has existing inspection records.', 'danger')
@@ -112,7 +118,9 @@ def delete_facility(facility_id):
 @login_required
 @supervisor_required
 def create_area(facility_id):
-    facility = Facility.query.get_or_404(facility_id)
+    facility = db.session.get(Facility, facility_id)
+    if facility is None:
+        abort(404)
     form = AreaForm()
     form.facility_id.choices = [(facility.id, facility.name)]
     form.facility_id.data = facility.id
@@ -137,7 +145,9 @@ def create_area(facility_id):
 @login_required
 @supervisor_required
 def edit_area(area_id):
-    area = Area.query.get_or_404(area_id)
+    area = db.session.get(Area, area_id)
+    if area is None:
+        abort(404)
     form = AreaForm(obj=area)
     
     facilities = Facility.query.filter_by(active=True).order_by(Facility.name).all()
@@ -160,7 +170,9 @@ def edit_area(area_id):
 @login_required
 @supervisor_required
 def delete_area(area_id):
-    area = Area.query.get_or_404(area_id)
+    area = db.session.get(Area, area_id)
+    if area is None:
+        abort(404)
     facility_id = area.facility_id
     
     if area.inspections.count() > 0: