Issue/86 move functions from ocf dashboard

pvsite_datamodel/read/__init__.py

@@ -12,4 +12,4 @@
     get_sites_from_user,
 )
 from .status import get_latest_status
-from .user import get_site_group_by_name, get_user_by_email
+from .user import get_all_site_groups, get_all_users, get_site_group_by_name, get_user_by_email

pvsite_datamodel/read/user.py

@@ -1,5 +1,6 @@
 """ Functions for reading user data from the database. """
 import logging
+from typing import List
 
 from sqlalchemy.orm import Session
 
@@ -35,6 +36,22 @@ def get_user_by_email(session: Session, email: str):
     return user
 
 
+# get all users
+def get_all_users(session: Session) -> List[UserSQL]:
+    """
+    Get all users from the database.
+
+    :param session: database session
+    """
+    query = session.query(UserSQL)
+
+    query = query.order_by(UserSQL.email.asc())
+
+    users = query.all()
+
+    return users
+
+
 def get_site_group_by_name(session: Session, site_group_name: str):
     """
     Get site group by name. If site group does not exist, make one.
@@ -57,3 +74,19 @@ def get_site_group_by_name(session: Session, site_group_name: str):
         session.commit()
 
     return site_group
+
+
+# get all site groups
+def get_all_site_groups(session: Session) -> List[SiteGroupSQL]:
+    """
+    Get all site groups from the database.
+
+    :param session: database session
+    """
+    query = session.query(SiteGroupSQL)
+
+    query = query.order_by(SiteGroupSQL.site_group_name.asc())
+
+    site_groups = query.all()
+
+    return site_groups

pvsite_datamodel/write/user_and_site.py

@@ -1,7 +1,11 @@
 """ Tools for making fake users and sites in the database."""
 from datetime import datetime, timezone
 
-from pvsite_datamodel.sqlmodels import SiteGroupSQL, SiteSQL, UserSQL
+import sqlalchemy as sa
+from sqlalchemy.orm.session import Session
+
+from pvsite_datamodel.read import get_user_by_email
+from pvsite_datamodel.sqlmodels import ForecastSQL, ForecastValueSQL, SiteGroupSQL, SiteSQL, UserSQL
 
 
 def make_site(db_session, ml_id=1):
@@ -50,3 +54,176 @@ def make_user(db_session, email, site_group):
     db_session.commit()
 
     return user
+
+
+# TODO move create_new site_function here.
+#  For the moment, data folder in ocf-dashboard has code this function depends on and this
+#  should probably be done in an new issue
+
+
+def create_user(
+    session: Session,
+    email: str,
+    site_group_name: str,
+) -> UserSQL:
+    """Create a user.
+
+    :param session: database session
+    :param email: email of user being created
+    :param site_group: name of the site group this user will be part of
+    """
+
+    site_group = (
+        session.query(SiteGroupSQL).filter(SiteGroupSQL.site_group_name == site_group_name).first()
+    )
+
+    user = UserSQL(email=email, site_group_uuid=site_group.site_group_uuid)
+
+    session.add(user)
+    site_group.users.append(user)
+    session.commit()
+
+    return user
+
+
+# update functions for site and site group
+def add_site_to_site_group(session: Session, site_uuid: str, site_group_name: str) -> SiteGroupSQL:
+    """Add a site to a site group.
+
+    NB: Sites can belong to many site groups.
+    :param session: database session
+    :param site_uuid: uuid of site
+    :param site_group_name: name of site group
+    """
+    site_group = (
+        session.query(SiteGroupSQL).filter(SiteGroupSQL.site_group_name == site_group_name).first()
+    )
+
+    site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).one()
+
+    if site not in site_group.sites:
+        site_group.sites.append(site)
+
+    session.commit()
+
+    return site_group.sites
+
+
+# change site group for user
+def change_user_site_group(session, email: str, site_group_name: str):
+    """
+    Change user to a specific site group name.
+
+    :param session: the database session
+    :param email: the email of the user
+    """
+    update_user_site_group(session=session, email=email, site_group_name=site_group_name)
+    user = get_user_by_email(session=session, email=email)
+    user_site_group = user.site_group.site_group_name
+    user = user.email
+    return user, user_site_group
+
+
+def update_user_site_group(session: Session, email: str, site_group_name: str) -> UserSQL:
+    """Change site group for user.
+
+    :param session: database session
+    :param email: email of user
+    :param site_group_name: name of site group
+    """
+    site_group = (
+        session.query(SiteGroupSQL).filter(SiteGroupSQL.site_group_name == site_group_name).first()
+    )
+
+    user = session.query(UserSQL).filter(UserSQL.email == email)
+
+    user = user.update({"site_group_uuid": site_group.site_group_uuid})
+
+    session.commit()
+
+
+# delete functions for site, user, and site group
+def delete_site(session: Session, site_uuid: str) -> SiteGroupSQL:
+    """Delete a site group.
+
+    :param session: database session
+    :param site_uuid: unique identifier for site
+    """
+    site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()
+
+    forecast_uuids = (
+        session.query(ForecastSQL.forecast_uuid).filter(ForecastSQL.site_uuid == site_uuid).all()
+    )
+
+    forecast_uuids = [str(forecast_uuid[0]) for forecast_uuid in forecast_uuids]
+
+    # get and delete all forecast values for site
+    stmt = sa.delete(ForecastValueSQL).where(ForecastValueSQL.forecast_uuid.in_(forecast_uuids))
+    session.execute(stmt)
+
+    stmt_2 = sa.delete(ForecastSQL).where(ForecastSQL.forecast_uuid.in_(forecast_uuids))
+    session.execute(stmt_2)
+
+    # we decided not to delete generation data for the site
+    # because it seems sensible to keep it for now
+
+    session.delete(site)
+
+    message = f"Site with site uuid {site.site_uuid} deleted successfully"
+
+    session.commit()
+
+    return message
+
+
+# delete user
+def delete_user(session: Session, email: str) -> UserSQL:
+    """Delete a user.
+
+    :param session: database session
+    :param email: email of user being deleted
+    """
+    user = session.query(UserSQL).filter(UserSQL.email == email).first()
+
+    session.delete(user)
+
+    message = (
+        f"User with email {user.email} and site_group_uuid "
+        f"{user.site_group_uuid} deleted successfully"
+    )
+
+    session.commit()
+
+    return message
+
+
+# delete site group
+def delete_site_group(session: Session, site_group_name: str) -> SiteGroupSQL:
+    """Delete a site group.
+
+    :param session: database session
+    :param site_group_name: name of site group being deleted
+    """
+    site_group = (
+        session.query(SiteGroupSQL).filter(SiteGroupSQL.site_group_name == site_group_name).first()
+    )
+
+    site_group_users = site_group.users
+
+    if len(site_group_users) > 0:
+        message = (
+            f"Site group with name {site_group.site_group_name} and "
+            f"site group uuid {site_group.site_group_uuid} has users and cannot be deleted."
+        )
+        return message
+
+    session.delete(site_group)
+
+    message = (
+        f"Site group with name {site_group.site_group_name} "
+        f"and site group uuid {site_group.site_group_uuid} deleted successfully."
+    )
+
+    session.commit()
+
+    return message

scripts/add_all_sites_to_user.py

@@ -28,27 +28,39 @@
 
     print('Found users: ' + str(len(users)))
 
-    for user in users:
-        print(user.email)
+    site_group = get_site_group_by_name(session=session, site_group_name='ocf')
 
-        # 1. attached to correct site group
-        if user.site_group.site_group_name != 'ocf':
-            site_group = get_site_group_by_name(session=session, site_group_name='ocf')
-            user.site_group = site_group
-
-            session.commit()
-
-    # # 2. make sure all sites are in ocf site group
-    # all_sites = get_all_sites(session=session)
+    # for user in users:
+    #     print(user.email)
     #
-    # site_uuids = [site.site_uuid for site in site_group.sites]
+    #     # 1. attached to correct site group
+    #     if user.site_group.site_group_name != 'ocf':
+    #         site_group = get_site_group_by_name(session=session, site_group_name='ocf')
+    #         user.site_group = site_group
     #
-    # for site in all_sites:
-    #     if site.site_uuid not in site_uuids:
-    #         print(f'Adding site {site.site_uuid} to user {user_email}')
-    #         site_group.sites.append(site)
-    #
-    # session.commit()
+    #         session.commit()
+
+    # 2. make sure all sites are in ocf site group
+    all_sites = get_all_sites(session=session)
+    print(f'Found {len(all_sites)} sites')
+
+    site_uuids = [site.site_uuid for site in site_group.sites]
+
+    for site in all_sites:
+        if site.site_uuid not in site_uuids:
+            print(f'Adding site {site.site_uuid} in ocf site group')
+            site_group.sites.append(site)
+        else:
+            print(f'Site {site.site_uuid} already in ocf site group')
+
+        # if site.client_site_name == '38 Norreys Avenue':
+        #     for f in site.forecasts:
+        #         for fv in f.forecast_values:
+        #             session.delete(fv)
+        #         session.delete(f)
+        #     session.delete(site)
+
+    session.commit()
 
 
 

tests/test_delete.py

@@ -0,0 +1,84 @@
+"""Test Delete Functions"""
+
+from pvsite_datamodel.sqlmodels import SiteSQL, UserSQL
+from pvsite_datamodel.write.user_and_site import (
+    delete_site,
+    delete_site_group,
+    delete_user,
+    make_site,
+    make_site_group,
+    make_user,
+)
+
+
+# test to delete a site
+def test_delete_site(db_session):
+    """Test to delete a site."""
+
+    site_1 = make_site(db_session=db_session, ml_id=1)
+
+    # todo add site to site group
+
+    site_group = make_site_group(db_session=db_session, site_group_name="test_site_group")
+    site_group.sites.append(site_1)
+
+    site_uuid = site_1.site_uuid
+
+    message = delete_site(session=db_session, site_uuid=site_uuid)
+
+    site = db_session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()
+
+    assert site is None
+    assert message == f"Site with site uuid {site_uuid} deleted successfully"
+
+
+def test_delete_user(db_session):
+    """Test to delete a user."""
+
+    site_group_1 = make_site_group(db_session=db_session)
+
+    user_1 = make_user(
+        db_session=db_session,
+        email="[email protected]",
+        site_group=site_group_1,
+    )
+
+    message = delete_user(session=db_session, email=user_1.email)
+
+    user = db_session.query(UserSQL).filter(UserSQL.email == user_1.email).first()
+
+    assert (
+        message == f"User with email {user_1.email} and "
+        f"site_group_uuid {user_1.site_group_uuid} deleted successfully"
+    )
+    assert user is None
+
+
+def test_delete_site_group(db_session):
+    site_group = make_site_group(db_session=db_session, site_group_name="test_site_group")
+
+    message = delete_site_group(session=db_session, site_group_name=site_group.site_group_name)
+
+    assert (
+        message == f"Site group with name {site_group.site_group_name} and "
+        f"site group uuid {site_group.site_group_uuid} deleted successfully."
+    )
+
+
+def test_delete_site_group_with_users(db_session):
+    """Test to delete a site group with users."""
+
+    site_group_1 = make_site_group(db_session=db_session)
+
+    _ = make_user(
+        db_session=db_session,
+        email="[email protected]",
+        site_group=site_group_1,
+    )
+
+    message = delete_site_group(session=db_session, site_group_name=site_group_1.site_group_name)
+
+    assert (
+        message == f"Site group with name {site_group_1.site_group_name} and "
+        f"site group uuid {site_group_1.site_group_uuid} has users and cannot be deleted."
+    )