Adding the clear_portfolio method for quickly emptying out a portfolio for easy rebuilds

backend/app/db/functions/recommendations_functions.py

@@ -1,10 +1,14 @@
-from sqlalchemy import insert
+from sqlalchemy import insert, delete
+from sqlalchemy.orm import Session
 from backend.app.db.models.recommendations import Plan, Recommendation, RecommendationMaterials, PlanRecommendations
+from backend.app.db.models.portfolio import PropertyModel, PropertyTargetsModel, PropertyDetailsMeter, \
+    PropertyDetailsEpcModel
 
 
-def create_plan(session, plan):
+def create_plan(session: Session, plan):
     """
     This function will create a record for the plan in the database if it does not exist.
+    :param session: The database session
     :param plan: dictionary of data representing a plan to be created
     """
 
@@ -15,7 +19,7 @@ def create_plan(session, plan):
     return new_plan.id
 
 
-def create_recommendation(session, recommendation):
+def create_recommendation(session: Session, recommendation):
     """
     This function will create a record for the recommendation in the database if it does not exist.
     :param session: The database session
@@ -29,7 +33,7 @@ def create_recommendation(session, recommendation):
     return new_recommendation.id
 
 
-def create_recommendation_material(session, recommendation_id, material_id, depth):
+def create_recommendation_material(session: Session, recommendation_id, material_id, depth):
     """
     This function will create a record for the recommendation_material in the database if it does not exist.
     :param session: The databse session
@@ -49,9 +53,10 @@ def create_recommendation_material(session, recommendation_id, material_id, dept
     return new_recommendation_material.id
 
 
-def create_plan_recommendations(session, plan_id, recommendation_ids):
+def create_plan_recommendations(session: Session, plan_id, recommendation_ids):
     """
     This function will create records for the plan_recommendation in the database.
+    :param session: The database session
     :param plan_id: ID of the plan
     :param recommendation_ids: list of recommendation IDs
     """
@@ -63,7 +68,7 @@ def create_plan_recommendations(session, plan_id, recommendation_ids):
     session.execute(insert(PlanRecommendations).values(data))
 
 
-def upload_recommendations(session, recommendations_to_upload, property_id):
+def upload_recommendations(session: Session, recommendations_to_upload, property_id):
     # Prepare data for bulk insert for Recommendation
     recommendations_data = [
         {
@@ -112,3 +117,39 @@ def upload_recommendations(session, recommendations_to_upload, property_id):
     session.flush()
 
     return uploaded_recommendation_ids
+
+
+def clear_portfolio(session: Session, portfolio_id: int):
+    # Fetch all property IDs associated with the given portfolio
+    property_ids = session.query(PropertyModel.id).filter(PropertyModel.portfolio_id == portfolio_id).all()
+    property_ids = [p.id for p in property_ids]
+
+    # Fetch all recommendation IDs associated with the properties
+    recommendation_ids = session.query(Recommendation.id).filter(Recommendation.property_id.in_(property_ids)).all()
+    recommendation_ids = [r.id for r in recommendation_ids]
+
+    # Delete all entries from RecommendationMaterials for these recommendations
+    session.execute(
+        delete(RecommendationMaterials).where(RecommendationMaterials.recommendation_id.in_(recommendation_ids))
+    )
+
+    # Delete all entries from PlanRecommendations that reference plans in the portfolio
+    session.execute(delete(PlanRecommendations).where(PlanRecommendations.plan_id.in_(
+        session.query(Plan.id).filter(Plan.portfolio_id == portfolio_id).subquery().as_scalar()
+    )))
+
+    # Delete all Plans associated with the portfolio
+    session.execute(delete(Plan).where(Plan.portfolio_id == portfolio_id))
+
+    # Delete all Recommendations associated with the properties
+    session.execute(delete(Recommendation).where(Recommendation.property_id.in_(property_ids)))
+
+    # Now, delete the PropertyModels and related details
+    # Delete PropertyTargetsModel, PropertyDetailsMeter, PropertyDetailsEpcModel, and PropertyModel
+    session.execute(delete(PropertyTargetsModel).where(PropertyTargetsModel.portfolio_id == portfolio_id))
+    # session.execute(delete(PropertyDetailsMeter).where(PropertyDetailsMeter.uprn.in_(property_ids)))
+    session.execute(delete(PropertyDetailsEpcModel).where(PropertyDetailsEpcModel.portfolio_id == portfolio_id))
+    session.execute(delete(PropertyModel).where(PropertyModel.portfolio_id == portfolio_id))
+
+    # Commit the changes
+    session.commit()

backend/app/plan/router.py

@@ -312,7 +312,7 @@ async def trigger_plan(body: PlanTriggerRequest):
 
         logger.info("Uploading recommendations to the database")
         session.commit()
-        for i in tqdm(range(0, len(input_properties), BATCH_SIZE)):
+        for i in range(0, len(input_properties), BATCH_SIZE):
             try:
                 # Take a slice of the input_properties list to make a batch
                 batch_properties = input_properties[i:i + BATCH_SIZE]