adding other energy assessment tables

This commit is contained in:
Khalim Conn-Kowlessar 2024-09-04 15:58:08 +01:00
parent cdb1bebddc
commit d8e337e55d
2 changed files with 77 additions and 23 deletions

View file

@ -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
}

View file

@ -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()