Merge pull request #14 from openclimatefix/add-read-all-locations

add method to read all lcations + tests

Author: peterdudfield

nowcasting_datamodel/read/read.py

@@ -175,6 +175,33 @@ def get_location(session: Session, gsp_id: int) -> LocationSQL:
     return location
 
 
+def get_all_location(session: Session, gsp_ids: List[int] = None) -> List[LocationSQL]:
+    """
+    Get all location object from gsp id
+
+    :param session: database session
+    :param gsp_ids: list of gsp id of the location
+
+    return: List of GSP locations
+
+    """
+
+    # start main query
+    query = session.query(LocationSQL)
+    query = query.distinct(LocationSQL.gsp_id)
+
+    # filter on gsp_id
+    if gsp_ids is not None:
+        query = query.filter(LocationSQL.gsp_id.in_(gsp_ids))
+
+    query = query.order_by(LocationSQL.gsp_id)
+
+    # get all results
+    locations = query.all()
+
+    return locations
+
+
 def get_model(session: Session, name: str, version: str) -> MLModelSQL:
     """
     Get model object from name and version

tests/test_read.py

@@ -8,6 +8,7 @@
 from nowcasting_datamodel.models import Forecast, ForecastValue, LocationSQL, MLModel, PVSystem
 from nowcasting_datamodel.read.read import (
     get_all_gsp_ids_latest_forecast,
+    get_all_location,
     get_forecast_values,
     get_latest_forecast,
     get_latest_national_forecast,
@@ -19,6 +20,15 @@
 logger = logging.getLogger(__name__)
 
 
+def test_get_all_location(db_session):
+
+    db_session.add(LocationSQL(label="GSP_1", gsp_id=1))
+    db_session.add(LocationSQL(label="GSP_2", gsp_id=2))
+
+    locations = get_all_location(session=db_session)
+    assert len(locations) == 2
+
+
 def test_get_model(db_session):
 
     model_read_1 = get_model(session=db_session, name="test_name", version="9.9.9")