From d8e337e55dde723f2f29766716b0b4e3233c066a Mon Sep 17 00:00:00 2001 From: Khalim Conn-Kowlessar Date: Wed, 4 Sep 2024 15:58:08 +0100 Subject: [PATCH] adding other energy assessment tables --- backend/app/db/models/energy_assessments.py | 27 +++++++- backend/app/energy_assessments/router.py | 73 ++++++++++++++------- 2 files changed, 77 insertions(+), 23 deletions(-) diff --git a/backend/app/db/models/energy_assessments.py b/backend/app/db/models/energy_assessments.py index 3928f9fa..a5667a0a 100644 --- a/backend/app/db/models/energy_assessments.py +++ b/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 + } diff --git a/backend/app/energy_assessments/router.py b/backend/app/energy_assessments/router.py index 1e4b44f0..f0577c25 100644 --- a/backend/app/energy_assessments/router.py +++ b/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( - { - "uprn": uprn, - "document_type": doc_type, - "document_location": doc_files - } - ) + + for doc in doc_files: + property_ea_docs.append( + { + "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()