Update app/services/class_service.py

app/services/class_service.py

@@ -4,8 +4,9 @@ from typing import List, Optional, Iterable, Tuple
 
 from sqlalchemy.orm import Session
 
-from models.class_ import Class, ClassEnrollment
+from models.class_ import Class, ClassEnrollment, ClassMembershipPlan
 from models.membership import Membership
+from models.membership_plan import MembershipPlan
 from schemas.class_ import (
     ClassCreate,
     ClassUpdate,
@@ -145,10 +146,27 @@ def create_class(db: Session, class_in: ClassCreate) -> Class:
         max_students=class_in.max_students,
         is_active=getattr(class_in, "is_active", True),
         coach_email=getattr(class_in, "coach_email", None),
+        location=getattr(class_in, "location", None),
     )
     db.add(db_class)
     db.commit()
     db.refresh(db_class)
+    
+    # Handle membership plan associations
+    membership_plan_ids = getattr(class_in, "membership_plan_ids", [])
+    if membership_plan_ids:
+        for plan_id in membership_plan_ids:
+            # Verify plan exists
+            plan = db.query(MembershipPlan).filter(MembershipPlan.id == plan_id).first()
+            if plan:
+                class_plan = ClassMembershipPlan(
+                    class_id=db_class.id,
+                    membership_plan_id=plan_id,
+                )
+                db.add(class_plan)
+        db.commit()
+    
+    db.refresh(db_class)
     return db_class
 
 
@@ -191,6 +209,26 @@ def update_class(db: Session, class_id: int, class_in: ClassUpdate) -> Optional[
         db_class.is_active = class_in.is_active
     if hasattr(class_in, "coach_email") and class_in.coach_email is not None:
         db_class.coach_email = class_in.coach_email
+    if hasattr(class_in, "location"):
+        db_class.location = class_in.location
+
+    # Handle membership plan associations
+    if hasattr(class_in, "membership_plan_ids") and class_in.membership_plan_ids is not None:
+        # Remove existing associations
+        db.query(ClassMembershipPlan).filter(
+            ClassMembershipPlan.class_id == class_id
+        ).delete()
+        
+        # Add new associations
+        for plan_id in class_in.membership_plan_ids:
+            # Verify plan exists
+            plan = db.query(MembershipPlan).filter(MembershipPlan.id == plan_id).first()
+            if plan:
+                class_plan = ClassMembershipPlan(
+                    class_id=class_id,
+                    membership_plan_id=plan_id,
+                )
+                db.add(class_plan)
 
     db.commit()
     db.refresh(db_class)