diff --git a/backend/app/db/base.py b/backend/app/db/base.py index 59be7030..fa2b68a5 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -1,3 +1,5 @@ -from sqlalchemy.orm import declarative_base +from sqlalchemy.orm import DeclarativeBase -Base = declarative_base() + +class Base(DeclarativeBase): + pass diff --git a/backend/app/db/models/recommendations.py b/backend/app/db/models/recommendations.py index 9352eeb2..27d03303 100644 --- a/backend/app/db/models/recommendations.py +++ b/backend/app/db/models/recommendations.py @@ -54,19 +54,47 @@ class Recommendation(Base): class RecommendationMaterials(Base): __tablename__ = "recommendation_materials" - id = Column(BigInteger, primary_key=True, autoincrement=True) - recommendation_id = Column( - BigInteger, ForeignKey("recommendation.id"), nullable=False + id: Mapped[int] = mapped_column( + BigInteger, primary_key=True, autoincrement=True ) - material_id = Column(BigInteger, ForeignKey(Material.id), nullable=False) - created_at = Column(TIMESTAMP, nullable=False, server_default=func.now()) - depth = Column(Float, nullable=False) - quantity = Column(Float, nullable=False) - quantity_unit = Column( + + recommendation_id: Mapped[int] = mapped_column( + BigInteger, + ForeignKey("recommendation.id"), + nullable=False, + ) + + material_id: Mapped[int] = mapped_column( + BigInteger, + ForeignKey(Material.id), + nullable=False, + ) + + created_at: Mapped[datetime] = mapped_column( + TIMESTAMP, + nullable=False, + server_default=func.now(), + ) + + depth: Mapped[float] = mapped_column( + Float, + nullable=False, + ) + + quantity: Mapped[float] = mapped_column( + Float, + nullable=False, + ) + + quantity_unit: Mapped[QuantityUnits] = mapped_column( Enum(QuantityUnits, values_callable=lambda x: [e.value for e in x]), nullable=False, ) - estimated_cost = Column(Float, nullable=False) + + estimated_cost: Mapped[float] = mapped_column( + Float, + nullable=False, + ) class PlanTypeEnum(enum.Enum): # TODO: move this to domain? diff --git a/backend/export/property_scenarios/db_functions.py b/backend/export/property_scenarios/db_functions.py index a27806c2..1527a989 100644 --- a/backend/export/property_scenarios/db_functions.py +++ b/backend/export/property_scenarios/db_functions.py @@ -1,6 +1,8 @@ -from typing import List, Any, Dict, Optional +from typing import List, Any, Dict, Optional, Tuple, Sequence import pandas as pd +from sqlalchemy import select from sqlalchemy.orm import Session +from sqlalchemy.engine import Row from collections import defaultdict from backend.app.db.models.recommendations import ( @@ -20,7 +22,7 @@ logger = setup_logger() class DbMethods: - def __init__(self, session: Session): + def __init__(self, session: Session) -> None: self.session = session def get_properties(self, portfolio_id: int) -> pd.DataFrame: @@ -29,28 +31,31 @@ class DbMethods: :param portfolio_id: :return: """ - query = ( - self.session.query(PropertyModel, PropertyDetailsEpcModel) + stmt = ( + select(PropertyModel, PropertyDetailsEpcModel) .join( PropertyDetailsEpcModel, PropertyModel.id == PropertyDetailsEpcModel.property_id, ) - .filter(PropertyModel.portfolio_id == portfolio_id) - .all() + .where(PropertyModel.portfolio_id == portfolio_id) ) - data = [ + rows: Sequence[Row[Tuple[PropertyModel, PropertyDetailsEpcModel]]] = ( + self.session.execute(stmt).all() + ) + + data: List[Dict[str, Any]] = [ { **{ - col.name: getattr(row.PropertyModel, col.name) - for col in PropertyModel.__table__.columns + col.name: getattr(property_model, col.name) + for col in PropertyModel.__table__.columns.values() }, **{ - col.name: getattr(row.PropertyDetailsEpcModel, col.name) - for col in PropertyDetailsEpcModel.__table__.columns + col.name: getattr(epc_model, col.name) + for col in PropertyDetailsEpcModel.__table__.columns.values() }, } - for row in query + for property_model, epc_model in rows ] return pd.DataFrame(data) @@ -92,11 +97,11 @@ class DbMethods: # DISTINCT ON (property_id) keeps the first row per property, # ordered by created_at DESC so we get the newest one. - plans_query = ( - self.session.query(PlanModel) - .filter( + stmt = ( + select(PlanModel) + .where( PlanModel.portfolio_id == portfolio_id, - PlanModel.is_default.is_(True) + PlanModel.is_default.is_(True), ) .distinct(PlanModel.property_id) .order_by( @@ -110,11 +115,13 @@ class DbMethods: # DISTINCT ON (scenario_id, property_id) keeps the newest # plan per scenario/property combination. - plans_query = ( - self.session.query(PlanModel) - .filter( + assert scenario_ids is not None + + stmt = ( + select(PlanModel) + .where( PlanModel.portfolio_id == portfolio_id, - PlanModel.scenario_id.in_(scenario_ids) + PlanModel.scenario_id.in_(scenario_ids), ) .distinct( PlanModel.scenario_id, @@ -128,13 +135,14 @@ class DbMethods: ) logger.info("Fetching plans") - plans = plans_query.all() + + plans: Sequence[PlanModel] = self.session.scalars(stmt).all() return pd.DataFrame( [ { col.name: getattr(plan, col.name) - for col in PlanModel.__table__.columns + for col in PlanModel.__table__.columns.values() } for plan in plans ] @@ -146,35 +154,34 @@ class DbMethods: logger.info("No plan ids provided") return pd.DataFrame() - recs_query = ( - self.session.query( - Recommendation, - PlanModel.scenario_id, - PlanModel.name - ) + stmt = ( + select(Recommendation, PlanModel.scenario_id, PlanModel.name) .join( PlanRecommendations, Recommendation.id == PlanRecommendations.recommendation_id, ) .join(PlanModel, PlanModel.id == PlanRecommendations.plan_id) - .filter( + .where( PlanRecommendations.plan_id.in_(plan_ids), Recommendation.default.is_(True), Recommendation.already_installed.is_(False), ) - .all() ) - data = [ + rows: Sequence[Tuple[Recommendation, Optional[int], Optional[str]]] = ( + self.session.execute(stmt).tuples().all() + ) + + data: List[Dict[str, Any]] = [ { **{ - col.name: getattr(r.Recommendation, col.name) - for col in Recommendation.__table__.columns + col.name: getattr(rec_model, col.name) + for col in Recommendation.__table__.columns.values() }, - "scenario_id": r.scenario_id, - "plan_name": r.name, + "scenario_id": scenario_id, + "plan_name": plan_name, } - for r in recs_query + for rec_model, scenario_id, plan_name in rows ] return pd.DataFrame(data) @@ -185,12 +192,14 @@ class DbMethods: recommendations_df["materials"] = [] return recommendations_df - rec_ids = recommendations_df["id"].tolist() + rec_ids: List[int] = [int(x) for x in recommendations_df["id"].tolist()] - materials_query = ( - self.session.query(RecommendationMaterials) - .filter(RecommendationMaterials.recommendation_id.in_(rec_ids)) - .all() + stmt = select(RecommendationMaterials).where( + RecommendationMaterials.recommendation_id.in_(rec_ids) + ) + + materials_query: Sequence[RecommendationMaterials] = ( + self.session.scalars(stmt).all() ) materials_map: Dict[int, List[Dict[str, Any]]] = defaultdict(list) @@ -206,8 +215,8 @@ class DbMethods: } ) - recommendations_df["materials"] = recommendations_df["id"].apply( - lambda x: materials_map.get(x, []) + recommendations_df["materials"] = recommendations_df["id"].astype(int).apply( + lambda x: materials_map.get(int(x), []) ) return recommendations_df