fixing typing issues

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

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?

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