Add client table

README.md

@@ -40,6 +40,7 @@ Classes specifying table schemas:
 - SiteSQL
 - SiteGroupSQL
 - StatusSQL
+- ClientSQL
 
 Database connection objects:
 - DatabaseConnection
@@ -54,11 +55,13 @@ Currently available functions accessible via `from pvsite_datamodel.read import
 - get_site_by_uuid
 - get_site_by_client_site_id
 - get_site_by_client_site_name
+- get_sites_by_client_name 
 - get_all_sites
 - get_sites_by_country
 - get_site_group_by_name
 - get_latest_status
 - get_latest_forecast_values_by_site
+- get_client_by_name
 
 
 ### Write package functions
@@ -77,6 +80,9 @@ Currently available write functions accessible via `from pvsite_datamodels.write
 - delete_user
 - delete_site_group
 - make_fake_site
+- create_client
+- edit_client
+- assign_site_to_client
 
 
 ## Install the dependencies (requires [poetry][poetry])
@@ -140,7 +146,14 @@ classDiagram
         + inverter_capacity_kw : Float
         + module_capacity_kw : Float
         + ml_id : Integer ≪ U ≫
+        + client_uuid : UUID ≪ FK ≫
     }
+
+    class ClientSQL{
+        + client_uuid : UUID ≪ PK ≫
+        + client_name : String(255)
+    }
+
     class GenerationSQL{
         + generation_uuid : UUID ≪ PK ≫
         + site_uuid : UUID ≪ FK ≫
@@ -197,6 +210,7 @@ classDiagram
     MLModelSQL "1" -- "N" ForecastValueSQL : forecasts
     SiteSQL "1" -- "N" InverterSQL : contains
     UserSQL "1" -- "N" APIRequestSQL : performs_request
+    ClientSQL "1" -- "N" SiteSQL : owns
     class Legend{
     UUID: Universally Unique Identifier
     PK: Primary Key

alembic/versions/31d501a0aa52_add_client_table.py

@@ -0,0 +1,41 @@
+"""Add client table
+
+Revision ID: 31d501a0aa52
+Revises: 2a6e6975cd72
+Create Date: 2024-09-13 11:54:12.743980
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '31d501a0aa52'
+down_revision = '2a6e6975cd72'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('clients',
+        sa.Column('client_uuid', sa.UUID(), nullable=False),
+        sa.Column('client_name', sa.String(length=255), nullable=False),
+        sa.Column('created_utc', sa.DateTime(), nullable=True),
+        sa.PrimaryKeyConstraint('client_uuid')
+    )
+    op.create_index(op.f('ix_clients_client_name'), 'clients', ['client_name'], unique=True)
+    op.add_column('sites', sa.Column('client_uuid', sa.UUID(), nullable=True, comment='The UUID of the client this site belongs to'))
+    op.create_index(op.f('ix_sites_client_uuid'), 'sites', ['client_uuid'], unique=False)
+    op.create_foreign_key(None, 'sites', 'clients', ['client_uuid'], ['client_uuid'])
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'sites', type_='foreignkey')
+    op.drop_index(op.f('ix_sites_client_uuid'), table_name='sites')
+    op.drop_column('sites', 'client_uuid')
+    op.drop_index(op.f('ix_clients_client_name'), table_name='clients')
+    op.drop_table('clients')
+    # ### end Alembic commands ###

pvsite_datamodel/__init__.py

@@ -3,6 +3,7 @@
 from .connection import DatabaseConnection
 from .sqlmodels import (
     APIRequestSQL,
+    ClientSQL,
     ForecastSQL,
     ForecastValueSQL,
     GenerationSQL,

pvsite_datamodel/pydantic_models.py

@@ -2,6 +2,7 @@
 
 from datetime import datetime
 from typing import Optional
+from uuid import UUID
 
 from pydantic import BaseModel, Field
 
@@ -63,3 +64,6 @@ class PVSiteEditMetadata(BaseModel):
     capacity_kw: Optional[float] = Field(None, description="The site's total capacity in kw", ge=0)
     dno: Optional[str] = Field(None, description="The site's DNO")
     gsp: Optional[str] = Field(None, description="The site's GSP")
+    client_uuid: Optional[UUID] = Field(
+        None, description="The UUID of the client this site belongs to"
+    )

pvsite_datamodel/read/__init__.py

@@ -2,6 +2,7 @@
 Functions for reading from the PVSite database
 """
 
+from .client import get_client_by_name
 from .generation import get_pv_generation_by_sites, get_pv_generation_by_user_uuids
 from .latest_forecast_values import get_latest_forecast_values_by_site
 from .model import get_or_create_model
@@ -10,6 +11,7 @@
     get_site_by_client_site_id,
     get_site_by_client_site_name,
     get_site_by_uuid,
+    get_sites_by_client_name,
     get_sites_by_country,
     get_sites_from_user,
 )

pvsite_datamodel/read/client.py

@@ -0,0 +1,37 @@
+""" Functions for reading user data from the database. """
+
+import logging
+
+from sqlalchemy.orm import Session
+
+from pvsite_datamodel.sqlmodels import ClientSQL
+
+logger = logging.getLogger(__name__)
+
+
+def get_client_by_name(
+    session: Session, name: str, make_new_client_if_none: bool = True
+) -> ClientSQL:
+    """
+    Get client by name. If client does not exist, make one.
+
+    :param session: database session
+    :param name: name of the client
+    :param make_new_client_if_none: make client with name if doesn't exist
+    :return: client object
+    """
+
+    client = session.query(ClientSQL).filter(ClientSQL.client_name == name).first()
+
+    if client is None:
+        if make_new_client_if_none is True:
+            logger.info(f"Client with name {name} not found, so making one")
+
+            # make a new client
+            client = ClientSQL(client_name=name)
+            session.add(client)
+            session.commit()
+        else:
+            raise Exception(f"Could not find client with name {name}")
+
+    return client

pvsite_datamodel/read/site.py

@@ -6,7 +6,7 @@
 from sqlalchemy.orm import Session
 
 from pvsite_datamodel.pydantic_models import LatitudeLongitudeLimits
-from pvsite_datamodel.sqlmodels import SiteGroupSiteSQL, SiteGroupSQL, SiteSQL, UserSQL
+from pvsite_datamodel.sqlmodels import ClientSQL, SiteGroupSiteSQL, SiteGroupSQL, SiteSQL, UserSQL
 
 logger = logging.getLogger(__name__)
 
@@ -170,3 +170,33 @@ def get_sites_from_user(session, user, lat_lon_limits: Optional[LatitudeLongitud
     else:
         sites = user.site_group.sites
     return sites
+
+
+def get_sites_by_client_name(session: Session, client_name: str) -> List[SiteSQL]:
+    """Get sites from client name.
+
+    :param session: database session
+    :param client_name: client name
+    :return: list of site objects
+    """
+    logger.debug(f"Getting {client_name}'s sites")
+
+    # start main query
+    query = session.query(SiteSQL)
+
+    # join the Client table
+    query = query.join(ClientSQL)
+
+    # order by uuuid
+    query = query.order_by(SiteSQL.site_uuid)
+
+    # select the sites related to the client name
+    query = query.filter(ClientSQL.client_name == client_name)
+
+    # get all results
+    sites: Optional[List[SiteSQL]] = query.all()
+
+    if len(sites) == 0:
+        raise Exception(f"Could not find sites from client {client_name}")
+
+    return sites

pvsite_datamodel/sqlmodels.py

@@ -168,6 +168,14 @@ class SiteSQL(Base, CreatedMixin):
         comment="Auto-incrementing integer ID of the site for use in ML training",
     )
 
+    client_uuid = sa.Column(
+        UUID(as_uuid=True),
+        sa.ForeignKey("clients.client_uuid"),
+        nullable=True,
+        index=True,
+        comment="The UUID of the client this site belongs to",
+    )
+
     forecasts: Mapped[List["ForecastSQL"]] = relationship("ForecastSQL", back_populates="site")
     generation: Mapped[List["GenerationSQL"]] = relationship("GenerationSQL")
     inverters: Mapped[List["InverterSQL"]] = relationship(
@@ -176,6 +184,21 @@ class SiteSQL(Base, CreatedMixin):
     site_groups: Mapped[List["SiteGroupSQL"]] = relationship(
         "SiteGroupSQL", secondary="site_group_sites", back_populates="sites"
     )
+    client: Mapped[List["ClientSQL"]] = relationship("ClientSQL", back_populates="sites")
+
+
+class ClientSQL(Base, CreatedMixin):
+    """Class representing the client table.
+
+    Each client row specifies a single client.
+    """
+
+    __tablename__ = "clients"
+
+    client_uuid = sa.Column(UUID(as_uuid=True), default=uuid.uuid4, primary_key=True)
+    client_name = sa.Column(sa.String(255), nullable=False, index=True, unique=True)
+
+    sites: Mapped[List[SiteSQL]] = relationship("SiteSQL", back_populates="client")
 
 
 class GenerationSQL(Base, CreatedMixin):

pvsite_datamodel/write/__init__.py

@@ -2,6 +2,7 @@
 Functions for writing to the PVSite database
 """
 
+from .client import assign_site_to_client, create_client, edit_client
 from .forecast import insert_forecast_values
 from .generation import insert_generation_values
 from .user_and_site import (
@@ -13,6 +14,7 @@
     delete_site,
     delete_site_group,
     delete_user,
+    edit_site,
     make_fake_site,
     update_user_site_group,
 )

pvsite_datamodel/write/client.py

@@ -0,0 +1,65 @@
+"""
+Tools for making clients in the database.
+"""
+
+import logging
+from uuid import UUID
+
+from sqlalchemy.orm.session import Session
+
+from pvsite_datamodel.sqlmodels import ClientSQL, SiteSQL
+
+_log = logging.getLogger(__name__)
+
+
+def create_client(session: Session, client_name: str) -> ClientSQL:
+    """Create a client.
+
+    :param session: database session
+    :param client_name: name of client being created
+    """
+    client = ClientSQL(client_name=client_name)
+
+    session.add(client)
+    session.commit()
+
+    return client
+
+
+def edit_client(session: Session, client_uuid: UUID, client_name: str) -> ClientSQL:
+    """Edit an existing client.
+
+    :param session: database session
+    :param client_uuid: the existing client uuid
+    :param client_name: name of the client
+    """
+    client = session.query(ClientSQL).filter(ClientSQL.client_uuid == client_uuid).first()
+
+    client.client_name = client_name
+
+    session.add(client)
+    session.commit()
+
+    return client
+
+
+def assign_site_to_client(session: Session, site_uuid: str, client_name: str) -> str:
+    """Assign site to client.
+
+    :param session: database session
+    :param site_uuid: uuid of site
+    :param client_name: name of the client the site will be assigned to.
+    """
+
+    client = session.query(ClientSQL).filter(ClientSQL.client_name == client_name).first()
+
+    site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()
+
+    site.client_uuid = client.client_uuid
+
+    session.add(site)
+    session.commit()
+
+    message = f"Site with site uuid {site_uuid} successfully assigned to the client {client_name}"
+
+    return message

pvsite_datamodel/write/user_and_site.py

@@ -3,6 +3,7 @@
 import json
 from datetime import datetime, timezone
 from typing import Optional, Tuple
+from uuid import UUID
 
 import sqlalchemy as sa
 from sqlalchemy.orm.session import Session
@@ -74,6 +75,7 @@ def create_site(
     tilt: Optional[float] = None,
     inverter_capacity_kw: Optional[float] = None,
     module_capacity_kw: Optional[float] = None,
+    client_uuid: Optional[UUID] = None,
     ml_id: Optional[int] = None,
 ) -> [SiteSQL, str]:
     """
@@ -150,6 +152,7 @@ def create_site(
         tilt=tilt,
         inverter_capacity_kw=inverter_capacity_kw,
         module_capacity_kw=module_capacity_kw,
+        client_uuid=client_uuid,
     )
 
     session.add(site)
@@ -269,6 +272,7 @@ def edit_site(
         - capacity_kw: capacity of site in kw
         - dno: dno of site
         - gsp: gsp of site
+        - client_uuid: The UUID of the client this site belongs to
     """
     site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()