Updated functionalities

app/routes/customers.py

@@ -13,7 +13,7 @@ Provides a single screen to:
 """
 
 import logging
-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.user import User
@@ -238,7 +238,9 @@ def _send_invite_email(user, token):
 @supervisor_required
 def resend_invite(customer_id):
     """Generate a fresh token and resend the set-password invitation email."""
-    customer = User.query.get_or_404(customer_id)
+    customer = db.session.get(User, customer_id)
+    if customer is None:
+        abort(404)
     if customer.role != 'customer':
         flash('This action is only for customer accounts.', 'warning')
         return redirect(url_for('customers.index'))
@@ -297,7 +299,9 @@ def set_password(token):
 @login_required
 @supervisor_required
 def edit(customer_id):
-    customer = User.query.get_or_404(customer_id)
+    customer = db.session.get(User, customer_id)
+    if customer is None:
+        abort(404)
     if customer.role != 'customer':
         flash('This page is only for customer accounts.', 'warning')
         return redirect(url_for('customers.index'))
@@ -329,7 +333,9 @@ def edit(customer_id):
 @supervisor_required
 def manage(customer_id):
     """Single-customer detail page: profile + all assignments."""
-    customer = User.query.get_or_404(customer_id)
+    customer = db.session.get(User, customer_id)
+    if customer is None:
+        abort(404)
     if customer.role != 'customer':
         flash('This page is only for customer accounts.', 'warning')
         return redirect(url_for('customers.index'))
@@ -365,7 +371,9 @@ def manage(customer_id):
 @login_required
 @supervisor_required
 def add_assignment(customer_id):
-    customer = User.query.get_or_404(customer_id)
+    customer = db.session.get(User, customer_id)
+    if customer is None:
+        abort(404)
     if customer.role != 'customer':
         flash('Assignments are only for customer accounts.', 'warning')
         return redirect(url_for('customers.index'))
@@ -377,7 +385,9 @@ def add_assignment(customer_id):
         flash('Please select a contract.', 'warning')
         return redirect(url_for('customers.manage', customer_id=customer_id))
 
-    project = Project.query.get_or_404(project_id)
+    project = db.session.get(Project, project_id)
+    if project is None:
+        abort(404)
 
     # Guard: duplicate assignment
     existing = CustomerAssignment.query.filter_by(
@@ -414,7 +424,9 @@ def add_assignment(customer_id):
 @login_required
 @supervisor_required
 def remove_assignment(assignment_id):
-    assignment  = CustomerAssignment.query.get_or_404(assignment_id)
+    assignment  = db.session.get(CustomerAssignment, assignment_id)
+    if assignment is None:
+        abort(404)
     customer_id = assignment.user_id
     customer    = db.session.get(User, customer_id)
     project     = db.session.get(Project, assignment.project_id)
@@ -439,7 +451,9 @@ def remove_assignment(assignment_id):
 @login_required
 @supervisor_required
 def toggle_active(customer_id):
-    customer = User.query.get_or_404(customer_id)
+    customer = db.session.get(User, customer_id)
+    if customer is None:
+        abort(404)
     if customer.role != 'customer':
         flash('This action is only for customer accounts.', 'warning')
         return redirect(url_for('customers.index'))
@@ -747,6 +761,8 @@ def bulk_import():
 @supervisor_required
 def facilities_for_project(project_id):
     from flask import jsonify
-    project    = Project.query.get_or_404(project_id)
+    project    = db.session.get(Project, project_id)
+    if project is None:
+        abort(404)
     facilities = project.facilities.filter_by(active=True).order_by(Facility.name).all()
     return jsonify([{'id': f.id, 'name': f.name} for f in facilities])