adding other energy assessment tables

backend/app/db/models/energy_assessments.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, Integer, BigInteger, Text, Float, DateTime, Boolean, Date
+from sqlalchemy import Column, Integer, BigInteger, Text, Float, DateTime, Boolean, Date, ForeignKey
 from sqlalchemy.ext.declarative import declarative_base
 
 Base = declarative_base()
@@ -163,3 +163,28 @@ class EnergyAssessment(Base):
     @staticmethod
     def empty_response():
         return {"epc": {}, "condition": {}}
+
+
+class EnergyAssessmentScenarios(Base):
+    __tablename__ = 'energy_assessment_scenarios'
+    id = Column(BigInteger, primary_key=True, autoincrement=True)
+    scenario_name = Column(Text, nullable=False)
+    energy_assessment_id = Column(BigInteger, ForeignKey('energy_assessments.id'), nullable=False)
+
+
+class EnergyAssessmentDocuments(Base):
+    __tablename__ = 'energy_assessment_documents'
+    id = Column(BigInteger, primary_key=True, autoincrement=True)
+    uprn = Column(BigInteger, nullable=False)
+    energy_assessment_id = Column(BigInteger, ForeignKey('energy_assessments.id'), nullable=False)
+    document_type = Column(Text, nullable=False)  # You can handle this using an enum if needed
+    document_location = Column(Text, nullable=False)
+    uploaded_at = Column(DateTime(timezone=True), nullable=False)
+    scenario_id = Column(BigInteger, ForeignKey('energy_assessment_scenarios.id'), nullable=True)
+
+    @staticmethod
+    def empty_response():
+        return {
+            "id": None, "uprn": None, "document_type": None, "document_location": None, "uploaded_at": None,
+            "scenario_id": None
+        }

backend/app/energy_assessments/router.py

@@ -1,28 +1,24 @@
+import os
+from io import BytesIO
+
 from fastapi import APIRouter, Depends
 from starlette.responses import Response
 
-from backend.app.db.connection import db_engine
-
 from backend.app.config import get_settings
 from backend.app.dependencies import validate_token
 from backend.app.energy_assessments.schemas import EnergyAssessmentUploadPayload
 
+from sqlalchemy.orm import sessionmaker
 from sqlalchemy.exc import IntegrityError, OperationalError
-from sqlalchemy.orm import sessionmaker
-
-from backend.app.db.functions.energy_assessment_functions import bulk_insert_energy_assessments
-from sqlalchemy.orm import sessionmaker
 from backend.app.db.connection import db_engine
+from backend.app.db.functions.energy_assessment_functions import bulk_insert_energy_assessments
+
+from etl.xml_survey_extraction.XmlParser import XmlParser
+
 from utils.s3 import (
     read_from_s3, list_files_and_subfolders_in_s3_folder, list_xmls_in_s3_folder, save_csv_to_s3,
     list_files_in_s3_folder
 )
-from utils.logger import setup_logger
-from etl.xml_survey_extraction.XmlParser import XmlParser
-import os
-import pandas as pd
-from io import BytesIO
-
 from utils.logger import setup_logger
 
 logger = setup_logger()
@@ -81,7 +77,7 @@ async def upload(body: EnergyAssessmentUploadPayload):
             # We now split this into the different types of files
             # EPR
             eprs = [
-                file for file in energy_assessment_files if "EPR.pdf" in file.split("/")[-1].replace(" ", "").lower()
+                file for file in energy_assessment_files if "epr.pdf" in file.split("/")[-1].replace(" ", "").lower()
             ]
             # Condition report
             condition_reports = [
@@ -100,6 +96,31 @@ async def upload(body: EnergyAssessmentUploadPayload):
             # Floor plans - these are just the jpgs
             floor_plans = [file for file in energy_assessment_files if file.endswith(".jpg")]
 
+            # We now retrieve scenarios
+            scenario_folders = list_files_and_subfolders_in_s3_folder(
+                bucket_name=get_settings().ENERGY_ASSESSMENTS_BUCKET,
+                folder_name=assessment
+            )
+
+            # filter folders that contain the word scenario
+            scenario_folders = [
+                folder for folder in scenario_folders if "scenario" in folder.rstrip("/").split("/")[-1].lower()
+            ]
+            scenario_site_notes = []
+            scenario_draft_epcs = []
+            for sf in scenario_folders:
+                scenario_files = list_files_in_s3_folder(
+                    bucket_name=get_settings().ENERGY_ASSESSMENTS_BUCKET,
+                    folder_name=sf
+                )
+                notes = [
+                    file for file in scenario_files if "sitenotes" in file.split("/")[-1].replace(" ", "").lower()
+                ]
+                # This should be the leftovers
+                draft_epc = [file for file in scenario_files if file not in notes]
+                scenario_site_notes.extend(notes)
+                scenario_draft_epcs.extend(draft_epc)
+
             uprn = int(assessment.rstrip("/").split("/")[-1])
             assessments_map[uprn] = {
                 "xmls": uploaded_xmls,
@@ -107,26 +128,29 @@ async def upload(body: EnergyAssessmentUploadPayload):
                 "Condition Report": condition_reports,
                 "Evidence Report": evidence_reports,
                 "Summary Information": summary_reports,
-                "Floor PLan": floor_plans
+                "Floor Plan": floor_plans,
+                "Scenario Site Notes": scenario_site_notes,
+                "Scenario Draft EPC": scenario_draft_epcs
             }
 
         logger.info("Extracted energy assessment data and storing file locations to database")
         xml_data_to_store = []
         energy_assessment_documents = []
         for uprn, files in assessments_map.items():
-
             # Create the rows of data to insert into the energy assessment documents
             property_ea_docs = []
             for doc_type, doc_files in files.items():
                 if doc_type == "xmls":
                     continue
-                property_ea_docs.append(
+
-                    {
+                for doc in doc_files:
-                        "uprn": uprn,
+                    property_ea_docs.append(
-                        "document_type": doc_type,
+                        {
-                        "document_location": doc_files
+                            "uprn": uprn,
-                    }
+                            "document_type": doc_type,
-                )
+                            "document_location": doc
+                        }
+                    )
             energy_assessment_documents.extend(property_ea_docs)
 
             xmls = files["xmls"]
@@ -158,6 +182,11 @@ async def upload(body: EnergyAssessmentUploadPayload):
 
             xml_data_to_store.append(extracted_data)
 
+        logger.info("Storing energy assessment xml data to database")
+        bulk_insert_energy_assessments(session, xml_data_to_store)
+
+        # TODO: Store energy_assessment_documents
+
     except IntegrityError:
         logger.error("Database integrity error occurred", exc_info=True)
         session.rollback()