From 46c7a7c0d1ccf525a784879668d1723230bccff8 Mon Sep 17 00:00:00 2001 From: Khalim Conn-Kowlessar Date: Mon, 19 Aug 2024 13:46:03 +0100 Subject: [PATCH] added create_sfr_portfolio --- .../projects/midlands_portfolio/app.py | 47 +++++++++++++++++-- 1 file changed, 44 insertions(+), 3 deletions(-) diff --git a/etl/ownership/projects/midlands_portfolio/app.py b/etl/ownership/projects/midlands_portfolio/app.py index 8cad3c3e..2165dc94 100644 --- a/etl/ownership/projects/midlands_portfolio/app.py +++ b/etl/ownership/projects/midlands_portfolio/app.py @@ -1,5 +1,7 @@ from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm.exc import NoResultFound from backend.app.db.connection import db_engine +from backend.app.db.models.portfolio import Portfolio, PortfolioUsers from etl.ownership.Ownership import Ownership from etl.ownership.config import OWNERS_WHO_CANT_SELL as EXCLUDED_OWNERS @@ -59,6 +61,48 @@ PROPERTY_VALUE_ESTIMATE = 200_000 PORTFOLIO_VALUE = 75_000_000 +def create_sfr_portfolio(project_name, user_ids): + session = sessionmaker(bind=db_engine)() + session.begin() + + # Check for an existing portfolio by name + try: + portfolio = session.query(Portfolio).filter_by(name=project_name).one() + except NoResultFound: + portfolio = None + + if portfolio: + # Fetch the associated users + existing_user_ids = { + pu.userId for pu in session.query(PortfolioUsers.userId).filter_by(portfolioId=portfolio.id) + } + + # Check if the specified user_ids match any existing associations + if existing_user_ids.intersection(set(user_ids)): + print("Portfolio already exists under this name, for specified users.") + else: + print("Portfolio already exists under this name, for different users.") + return None # Optional: You could also update the user associations here if needed + + return portfolio # Return the existing portfolio data + + # If portfolio does not exist, create a new one + new_portfolio = Portfolio(name=project_name) + session.add(new_portfolio) + session.flush() # Ensures that 'id' is available before committing if needed + + # Create new user associations in PortfolioUsers + for user_id in user_ids: + new_association = PortfolioUsers(userId=user_id, portfolioId=new_portfolio.id) + session.add(new_association) + + session.commit() + print(f"New portfolio created with ID: {new_portfolio.id}") + session.close() + + return new_portfolio + + def app(): epc_column_filters = { "CURRENT_ENERGY_RATING": ["F", "G"] @@ -77,7 +121,4 @@ def app(): ) ownership_instance.pipeline(column_filters=epc_column_filters) - session = sessionmaker(bind=db_engine)() - session.begin() - # Create the project, if a portfolio doesn't exist for the project name