From b1f4f154ddb9371faa7cd49e9fbc52f02963bcbc Mon Sep 17 00:00:00 2001 From: Khalim Conn-Kowlessar Date: Tue, 30 Jul 2024 20:00:32 +0100 Subject: [PATCH] Refactored recommendation uploading to return ids explicitly on upload --- .../app/db/functions/portfolio_functions.py | 30 +++-- .../db/functions/recommendations_functions.py | 104 +++++++++--------- backend/app/db/models/recommendations.py | 1 + backend/app/plan/router.py | 10 +- backend/app/plan/schemas.py | 3 +- etl/xml_survey_extraction/app.py | 8 +- 6 files changed, 84 insertions(+), 72 deletions(-) diff --git a/backend/app/db/functions/portfolio_functions.py b/backend/app/db/functions/portfolio_functions.py index 008c4b8b..ffdabfb6 100644 --- a/backend/app/db/functions/portfolio_functions.py +++ b/backend/app/db/functions/portfolio_functions.py @@ -1,10 +1,14 @@ from sqlalchemy import func -from backend.app.db.models.recommendations import Plan, PlanRecommendations, Recommendation -from backend.app.db.models.portfolio import Portfolio +from backend.app.db.models.recommendations import Plan, PlanRecommendations, Recommendation, Scenario def aggregate_portfolio_recommendations( - session, portfolio_id: int, total_valuation_increase: float, labour_days: float, aggregated_data: dict + session, + portfolio_id: int, + scenario_id: int, + total_valuation_increase: float, + labour_days: float, + aggregated_data: dict ): # Aggregate multiple fields aggregates = ( @@ -17,7 +21,11 @@ def aggregate_portfolio_recommendations( ) .join(PlanRecommendations, PlanRecommendations.recommendation_id == Recommendation.id) .join(Plan, Plan.id == PlanRecommendations.plan_id) - .filter(Plan.portfolio_id == portfolio_id, Plan.is_default == True, Recommendation.default == True) + .filter( + Plan.portfolio_id == portfolio_id, + Plan.scenario_id == scenario_id, + Recommendation.default == True + ) .one() ) @@ -30,16 +38,16 @@ def aggregate_portfolio_recommendations( **aggregated_data } - # Get the portfolio and update the fields. This data needs to be stored against the plan, not the portfolio - portfolio = session.query(Portfolio).filter_by(id=portfolio_id).one() + # Get the scenario and update the fields. This data needs to be stored against the scenario, not the portfolio + portfolio_scenario = session.query(Scenario).filter_by(id=scenario_id).one() # Update the data for key, value in aggregates_dict.items(): - setattr(portfolio, key, value) + setattr(portfolio_scenario, key, value) # Insert total valuation increase and labour days - portfolio.property_valuation_increase = total_valuation_increase - portfolio.labour_days = labour_days + portfolio_scenario.property_valuation_increase = total_valuation_increase + portfolio_scenario.labour_days = labour_days - # Merge the updated portfolio back into the session - session.merge(portfolio) + # Merge the updated portfolio plan back into the session + session.merge(portfolio_scenario) session.flush() diff --git a/backend/app/db/functions/recommendations_functions.py b/backend/app/db/functions/recommendations_functions.py index c7765039..7ff09f22 100644 --- a/backend/app/db/functions/recommendations_functions.py +++ b/backend/app/db/functions/recommendations_functions.py @@ -95,62 +95,68 @@ def create_plan_recommendations(session: Session, plan_id, recommendation_ids): session.execute(insert(PlanRecommendations).values(data)) -def upload_recommendations(session: Session, recommendations_to_upload, property_id): - # Prepare data for bulk insert for Recommendation - recommendations_data = [ - { - "property_id": property_id, - "type": rec["type"], - "description": rec["description"], - "estimated_cost": rec["total"], - "default": rec["default"], - "starting_u_value": rec.get("starting_u_value"), - "new_u_value": rec.get("new_u_value"), - "sap_points": rec["sap_points"], - "energy_savings": rec["heat_demand"], - "kwh_savings": rec["kwh_savings"], - "co2_equivalent_savings": rec["co2_equivalent_savings"], - "total_work_hours": rec["labour_hours"], - "energy_cost_savings": rec["energy_cost_savings"], - "labour_days": rec["labour_days"], - "already_installed": rec["already_installed"], - } - for rec in recommendations_to_upload - ] +def upload_recommendations(session: Session, recommendations_to_upload, property_id, new_plan_id): + try: + # Prepare data for bulk insert for Recommendation + recommendations_data = [ + { + "property_id": property_id, + "type": rec["type"], + "description": rec["description"], + "estimated_cost": rec["total"], + "default": rec["default"], + "starting_u_value": rec.get("starting_u_value"), + "new_u_value": rec.get("new_u_value"), + "sap_points": rec["sap_points"], + "energy_savings": rec["heat_demand"], + "kwh_savings": rec["kwh_savings"], + "co2_equivalent_savings": rec["co2_equivalent_savings"], + "total_work_hours": rec["labour_hours"], + "energy_cost_savings": rec["energy_cost_savings"], + "labour_days": rec["labour_days"], + "already_installed": rec["already_installed"], + } + for rec in recommendations_to_upload + ] - session.bulk_insert_mappings(Recommendation, recommendations_data) + # Insert the recommendations, get back the IDs + stmt = insert(Recommendation).returning(Recommendation.id).values(recommendations_data) + result = session.execute(stmt) + uploaded_recommendation_ids = [row[0] for row in result] - # To get the IDs of the newly inserted recommendations, we need to flush the session - session.flush() + # Prepare data for bulk insert for RecommendationMaterials + recommendation_materials_data = [ + { + "recommendation_id": recommendation_id, + "material_id": part["id"], + "depth": int(part["depth"]) if part["depth"] else None, + "quantity": part["quantity"], + "quantity_unit": part["quantity_unit"], + "estimated_cost": part["total"], + } + for rec, recommendation_id in zip(recommendations_to_upload, uploaded_recommendation_ids) + for part in rec["parts"] + ] - # Map the uploaded_recommendation_ids with the original data for reference - uploaded_recommendation_ids = [rec.id for rec in session.query(Recommendation).filter( - Recommendation.property_id == property_id, - Recommendation.description.in_([rec["description"] for rec in recommendations_to_upload]) - )] + session.bulk_insert_mappings(RecommendationMaterials, recommendation_materials_data) - # Prepare data for bulk insert for RecommendationMaterials - # We can have multiple materials per recommendation. The aggregation of the materials will total the - # recommendation figures - recommendation_materials_data = [ - { - "recommendation_id": recommendation_id, - "material_id": part["id"], - "depth": int(part["depth"]) if part["depth"] else None, - "quantity": part["quantity"], - "quantity_unit": part["quantity_unit"], - "estimated_cost": part["total"], - } - for rec, recommendation_id in zip(recommendations_to_upload, uploaded_recommendation_ids) - for part in rec["parts"] - ] + # flush the changes to get the newly created IDs + session.flush() - session.bulk_insert_mappings(RecommendationMaterials, recommendation_materials_data) + create_plan_recommendations( + session, plan_id=new_plan_id, recommendation_ids=uploaded_recommendation_ids + ) - # flush the changes to get the newly created IDs - session.flush() + # Commit the transaction + session.commit() - return uploaded_recommendation_ids + return True + + except SQLAlchemyError as e: + # Rollback the transaction in case of an error + session.rollback() + print(f"An error occurred: {e}") + return False def clear_portfolio(session: Session, portfolio_id: int): diff --git a/backend/app/db/models/recommendations.py b/backend/app/db/models/recommendations.py index 6eddae1f..6ccfe7f7 100644 --- a/backend/app/db/models/recommendations.py +++ b/backend/app/db/models/recommendations.py @@ -53,6 +53,7 @@ class Plan(Base): name = Column(String, nullable=True, default="") portfolio_id = Column(BigInteger, ForeignKey(Portfolio.id), nullable=False) property_id = Column(BigInteger, ForeignKey(PropertyModel.id), nullable=False) + scenario_id = Column(BigInteger, ForeignKey('scenario.id')) # Doesn't have to be linked to a scenario created_at = Column(TIMESTAMP, nullable=False, server_default=func.now()) is_default = Column(Boolean, nullable=False) valuation_increase_lower_bound = Column(Float) diff --git a/backend/app/plan/router.py b/backend/app/plan/router.py index 4d73778e..a0d4e585 100644 --- a/backend/app/plan/router.py +++ b/backend/app/plan/router.py @@ -772,6 +772,7 @@ async def trigger_plan(body: PlanTriggerRequest): new_plan_id = create_plan(session, { "portfolio_id": body.portfolio_id, "property_id": p.id, + "scenario_id": engine_scenario.id, "is_default": True if p.is_new else False, "name": body.scenario_name, "valuation_increase_lower_bound": ( @@ -785,10 +786,8 @@ async def trigger_plan(body: PlanTriggerRequest): ), }) - uploaded_recommendation_ids = upload_recommendations(session, recommendations_to_upload, p.id) - - create_plan_recommendations( - session, plan_id=new_plan_id, recommendation_ids=uploaded_recommendation_ids + upload_recommendations( + session, recommendations_to_upload, p.id, new_plan_id ) property_valuation_increases.append( @@ -827,8 +826,7 @@ async def trigger_plan(body: PlanTriggerRequest): aggregate_portfolio_recommendations( session, portfolio_id=body.portfolio_id, - multi_plan=body.multi_plan, - + scenario_id=engine_scenario.id, total_valuation_increase=total_valuation_increase, labour_days=labour_days, aggregated_data=aggregated_data diff --git a/backend/app/plan/schemas.py b/backend/app/plan/schemas.py index b1e3a43a..108eb1ae 100644 --- a/backend/app/plan/schemas.py +++ b/backend/app/plan/schemas.py @@ -4,7 +4,6 @@ from typing import Optional class PlanTriggerRequest(BaseModel): budget: Optional[float] = None - # This can only have a fixed set of values goal: str housing_type: str goal_value: str @@ -36,7 +35,7 @@ class PlanTriggerRequest(BaseModel): "air_source_heat_pump", } - _allowed_goals = {"Increase EPC"} + _allowed_goals = {"Increasing EPC"} _allowed_housing_types = {"Social", "Private"} diff --git a/etl/xml_survey_extraction/app.py b/etl/xml_survey_extraction/app.py index ed2d20b6..a8bffc73 100644 --- a/etl/xml_survey_extraction/app.py +++ b/etl/xml_survey_extraction/app.py @@ -44,7 +44,7 @@ SCENARIOS = { "non_invasive_recommendations_file_path": "", "exclusions": ["floor_insulation", "fireplace"], "budget": None, - "Scenario Name": "Deep Retrofit", + "scenario_name": "Deep Retrofit", "multi_plan": True, }, # Scenario C, CWI, floor insulation, PV, AHSP @@ -59,7 +59,7 @@ SCENARIOS = { "non_invasive_recommendations_file_path": "", "exclusions": ["fireplace"], "budget": None, - "Scenario Name": "Whole House Retrofit", + "scenario_name": "Whole House Retrofit", "multi_plan": True, } ] @@ -80,7 +80,7 @@ SCENARIOS = { "non_invasive_recommendations_file_path": "", "exclusions": ["floor_insulation", "fireplace"], "budget": None, - "Scenario Name": "Deep Retrofit", + "scenario_name": "Deep Retrofit", "multi_plan": True, }, # Scenario B, floor insulation, PV, AHSP @@ -95,7 +95,7 @@ SCENARIOS = { "non_invasive_recommendations_file_path": "", "exclusions": ["fireplace"], "budget": None, - "Scenario Name": "Whole House Retrofit", + "scenario_name": "Whole House Retrofit", "multi_plan": True, } ]