Skip to content

Commit f2f4b42

Browse files
devsjcdevsjc
committed
chore(repo): Move to grpc library
1 parent 2a73beb commit f2f4b42

13 files changed

Lines changed: 155 additions & 148 deletions

File tree

‎pyproject.toml‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ dependencies = [
1818
"plotly==5.24.1",
1919
"psycopg2-binary==2.9.10",
2020
"SQLAlchemy==2.0.36",
21-
"streamlit==1.51.0",
21+
"streamlit==1.57.0",
2222
"uvicorn==0.34.0",
2323
"geopandas==1.0.1",
2424
"garrett-streamlit-auth0==0.9",
@@ -37,9 +37,7 @@ dependencies = [
3737
"matplotlib>=3.8,<4.0",
3838
"dp-sdk",
3939
"aiocache",
40-
"grpcio>=1.80.0",
41-
"grpcio-tools>=1.50.0",
42-
"grpc-requests>=0.1.17",
40+
"grpclib==0.4.8",
4341
"betterproto>=2.0.0b7",
4442
"pytest-asyncio>=1.3.0",
4543
"testcontainers>=4.14.0",

‎src/dataplatform/forecast/backend.py‎

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -6,29 +6,30 @@
66
import pandas as pd
77
import streamlit as st
88

9-
from ocf import dp
9+
from ocf.dp.dp_data import messages_pb2, service_pb2_grpc
10+
from ocf.dp.dp import common_pb2
1011

1112

1213
async def fetch_timeseries(
13-
client: dp.DataPlatformDataServiceStub,
14+
client: service_pb2_grpc.DataPlatformDataServiceStub,
1415
location_uuid: str,
1516
start_date: datetime.datetime,
1617
end_date: datetime.datetime,
1718
horizon_mins: int,
18-
forecasters: list[dp.Forecaster],
19+
forecasters: list[messages_pb2.Forecaster],
1920
init_times_utc: list[datetime.datetime] | None = None,
2021
) -> pd.DataFrame:
2122
"""Directly calls GetForecastAsTimeseries for selected models and init times."""
2223

23-
time_window = dp.TimeWindow(
24+
time_window = messages_pb2.TimeWindow(
2425
start_timestamp_utc=start_date, end_timestamp_utc=end_date
2526
)
2627
time_windows = []
2728
current_start = start_date
2829
while current_start < end_date:
2930
current_end = min(current_start + datetime.timedelta(days=7), end_date)
3031
time_windows.append(
31-
dp.TimeWindow(
32+
messages_pb2.TimeWindow(
3233
start_timestamp_utc=current_start,
3334
end_timestamp_utc=current_end
3435
)
@@ -38,32 +39,32 @@ async def fetch_timeseries(
3839
times_to_fetch = init_times_utc if init_times_utc else [None]
3940

4041
async def fetch_one(
41-
forecaster_obj: dp.Forecaster,
42-
window: dp.TimeWindow,
42+
forecaster_obj: messages_pb2.Forecaster,
43+
window: messages_pb2.TimeWindow,
4344
init_time: datetime.datetime | None
4445
):
45-
req = dp.GetForecastAsTimeseriesRequest(
46+
req = messages_pb2.GetForecastAsTimeseriesRequest(
4647
location_uuid=location_uuid,
47-
energy_source=dp.EnergySource.SOLAR,
48+
energy_source=common_pb2.EnergySource.ENERGY_SOURCE_SOLAR,
4849
horizon_mins=horizon_mins,
4950
time_window=window,
5051
forecaster=forecaster_obj,
5152
initialization_timestamp_utc=init_time,
5253
)
5354

5455
try:
55-
resp = await client.get_forecast_as_timeseries(req)
56+
resp = await client.GetForecastAsTimeseries(req)
5657
rows = []
5758
for val in resp.values:
5859
row = {
59-
"target_timestamp_utc": val.target_timestamp_utc,
60-
"initialization_timestamp_utc": val.initialization_timestamp_utc,
61-
"created_timestamp_utc": val.created_timestamp_utc,
60+
"target_timestamp_utc": val.target_timestamp_utc.ToDatetime(tzinfo=datetime.UTC),
61+
"initialization_timestamp_utc": val.initialization_timestamp_utc.ToDatetime(tzinfo=datetime.UTC),
62+
"created_timestamp_utc": val.created_timestamp_utc.ToDatetime(tzinfo=datetime.UTC),
6263
"effective_capacity_watts": val.effective_capacity_watts,
6364
"forecaster_name": forecaster_obj.forecaster_name,
6465
"location_uuid": resp.location_uuid,
6566
"horizon_mins": (
66-
val.target_timestamp_utc - val.initialization_timestamp_utc
67+
val.target_timestamp_utc.ToDatetime(tzinfo=datetime.UTC) - val.initialization_timestamp_utc.ToDatetime(tzinfo=datetime.UTC)
6768
).total_seconds()
6869
// 60,
6970
"p50_watts": int(
@@ -108,24 +109,24 @@ async def fetch_one(
108109

109110

110111
async def fetch_observations(
111-
client: dp.DataPlatformDataServiceStub,
112+
client: service_pb2_grpc.DataPlatformDataServiceStub,
112113
location_uuid: str,
113114
start_date: datetime.datetime,
114115
end_date: datetime.datetime,
115116
observers: list[str],
116-
energy_source: dp.EnergySource = dp.EnergySource.SOLAR,
117+
energy_source: common_pb2.EnergySource = common_pb2.EnergySource.ENERGY_SOURCE_SOLAR,
117118
) -> pd.DataFrame:
118119
"""Directly calls GetObservationsAsTimeseries for selected observers."""
119120

120-
time_window = dp.TimeWindow(
121+
time_window = messages_pb2.TimeWindow(
121122
start_timestamp_utc=start_date, end_timestamp_utc=end_date
122123
)
123124
time_windows = []
124125
current_start = start_date
125126
while current_start < end_date:
126127
current_end = min(current_start + datetime.timedelta(days=7), end_date)
127128
time_windows.append(
128-
dp.TimeWindow(
129+
messages_pb2.TimeWindow(
129130
start_timestamp_utc=current_start,
130131
end_timestamp_utc=current_end
131132
)
@@ -134,21 +135,21 @@ async def fetch_observations(
134135

135136

136137
# Run requests concurrently for all selected observers
137-
async def fetch_one(obs_name: str, window: dp.TimeWindow):
138-
req = dp.GetObservationsAsTimeseriesRequest(
138+
async def fetch_one(obs_name: str, window: messages_pb2.TimeWindow):
139+
req = messages_pb2.GetObservationsAsTimeseriesRequest(
139140
location_uuid=location_uuid,
140141
observer_name=obs_name,
141142
energy_source=energy_source,
142143
time_window=window,
143144
)
144145

145146
try:
146-
resp = await client.get_observations_as_timeseries(req)
147+
resp = await client.GetObservationsAsTimeseries(req)
147148
rows = []
148149
for val in resp.values:
149150
rows.append(
150151
{
151-
"target_timestamp_utc": val.timestamp_utc,
152+
"target_timestamp_utc": val.timestamp_utc.ToDatetime(tzinfo=datetime.UTC),
152153
"value_fraction": val.value_fraction,
153154
"effective_capacity_watts": val.effective_capacity_watts,
154155
"observer_name": obs_name,

‎src/dataplatform/forecast/cache.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from datetime import UTC, datetime, timedelta
44

5-
from ocf import dp
5+
from ocf.dp.dp_data import service_pb2_grpc
66

77
from dataplatform.forecast.constant import cache_seconds
88

@@ -11,7 +11,7 @@ def key_builder_remove_client(func: callable, *args: list, **kwargs: dict) -> st
1111
"""Custom key builder that ignores the client argument for caching purposes."""
1212
key = f"{func.__name__}:"
1313
for arg in args:
14-
if not isinstance(arg, dp.DataPlatformDataServiceStub):
14+
if not isinstance(arg, service_pb2_grpc.DataPlatformDataServiceStub):
1515
key += f"{arg}-"
1616

1717
for k, v in kwargs.items():

‎src/dataplatform/forecast/data.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,10 @@
1616

1717
async def get_forecast_data(
1818
dpc: service_pb2_grpc.DataPlatformDataServiceStub,
19-
location: dp.ListLocationsResponseLocationSummary,
19+
location: messages_pb2.ListLocationsResponse.LocationSummary,
2020
start_date: datetime,
2121
end_date: datetime,
22-
selected_forecasters: list[dp.Forecaster],
22+
selected_forecasters: list[messages_pb2.Forecaster],
2323
) -> pd.DataFrame:
2424
"""Get forecast data for the given location and time window."""
2525
all_data_df = []
@@ -66,10 +66,10 @@ async def get_forecast_data(
6666
@cached(ttl=cache_seconds, cache=Cache.MEMORY, key_builder=key_builder_remove_client)
6767
async def get_forecast_data_one_forecaster(
6868
dpc: service_pb2_grpc.DataPlatformDataServiceStub,
69-
location: dp.ListLocationsResponseLocationSummary,
69+
location: messages_pb2.ListLocationsResponse.LocationSummary,
7070
start_date: datetime,
7171
end_date: datetime,
72-
selected_forecaster: dp.Forecaster,
72+
selected_forecaster: messages_pb2.Forecaster,
7373
) -> pd.DataFrame | None:
7474
"""Get forecast data for one forecaster for the given location and time window."""
7575
all_data_list_dict = []
@@ -149,7 +149,7 @@ async def get_forecast_data_one_forecaster(
149149
@cached(ttl=cache_seconds, cache=Cache.MEMORY, key_builder=key_builder_remove_client)
150150
async def get_all_observations(
151151
client: service_pb2_grpc.DataPlatformDataServiceStub,
152-
location: dp.ListLocationsResponseLocationSummary,
152+
location: messages_pb2.ListLocationsResponse.LocationSummary,
153153
start_date: datetime,
154154
end_date: datetime,
155155
) -> pd.DataFrame:
@@ -225,10 +225,10 @@ async def get_all_observations(
225225

226226
async def get_all_data(
227227
client: service_pb2_grpc.DataPlatformDataServiceStub,
228-
selected_location: dp.ListLocationsResponseLocationSummary,
228+
selected_location: messages_pb2.ListLocationsResponse.LocationSummary,
229229
start_date: datetime,
230230
end_date: datetime,
231-
selected_forecasters: list[dp.Forecaster],
231+
selected_forecasters: list[messages_pb2.Forecaster],
232232
) -> dict:
233233
"""Get all forecast and observation data, and merge them."""
234234
# get generation data

‎src/dataplatform/forecast/main.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77

88
import pandas as pd
99
import streamlit as st
10-
from grpclib.client import Channel
11-
12-
from ocf import dp
10+
import grpc.aio
11+
from ocf.dp.dp import common_pb2
12+
from ocf.dp.dp_data import messages_pb2, service_pb2_grpc
1313

1414
from dataplatform.forecast.constant import metrics, observer_names
1515
from dataplatform.forecast.backend import fetch_observations, fetch_timeseries
@@ -49,8 +49,8 @@ async def async_dp_forecast_page() -> None:
4949
st.title("Data Platform Forecast Page")
5050
st.write("This is the forecast page from the Data Platform module.")
5151

52-
async with Channel(host=data_platform_host, port=data_platform_port) as channel:
53-
client = dp.DataPlatformDataServiceStub(channel)
52+
async with grpc.aio.insecure_channel(f"{data_platform_host}:{data_platform_port}") as channel:
53+
client = service_pb2_grpc.DataPlatformDataServiceStub(channel)
5454

5555
cfg = await setup_page(client)
5656
st.divider()
@@ -76,7 +76,7 @@ async def async_dp_forecast_page() -> None:
7676
start_date=cfg.start_date,
7777
end_date=cfg.end_date,
7878
observers=observer_names,
79-
energy_source=dp.EnergySource.SOLAR,
79+
energy_source=common_pb2.EnergySource.ENERGY_SOURCE_SOLAR,
8080
)
8181

8282
fetch_duration = (datetime.datetime.now() - start_time).total_seconds()

‎src/dataplatform/forecast/setup.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,20 @@
66
import pandas as pd
77
import streamlit as st
88
from aiocache import Cache, cached
9-
from ocf import dp
9+
from ocf.dp.dp import common_pb2
10+
from ocf.dp.dp_data import messages_pb2, service_pb2_grpc
1011

1112
from dataplatform.forecast.cache import key_builder_remove_client
1213
from dataplatform.forecast.constant import cache_seconds, metrics
1314

1415

1516
@cached(ttl=cache_seconds, cache=Cache.MEMORY, key_builder=key_builder_remove_client)
1617
async def get_location_names(
17-
client: dp.DataPlatformDataServiceStub,
18+
client: service_pb2_grpc.DataPlatformDataServiceStub,
1819
) -> dict:
1920
"""Get location names."""
20-
list_locations_request = dp.ListLocationsRequest()
21-
list_locations_response = await client.list_locations(list_locations_request)
21+
list_locations_request = messages_pb2.ListLocationsRequest()
22+
list_locations_response = await client.ListLocations(list_locations_request)
2223
all_locations = list_locations_response.locations
2324

2425
location_names = {loc.location_name: loc for loc in all_locations}
@@ -38,19 +39,19 @@ async def get_location_names(
3839

3940
@cached(ttl=cache_seconds, cache=Cache.MEMORY, key_builder=key_builder_remove_client)
4041
async def get_forecasters(
41-
client: dp.DataPlatformDataServiceStub,
42-
) -> list[dp.Forecaster]:
42+
client: service_pb2_grpc.DataPlatformDataServiceStub,
43+
) -> list[messages_pb2.Forecaster]:
4344
"""Get all forecasters."""
44-
get_forecasters_request = dp.ListForecastersRequest()
45-
get_forecasters_response = await client.list_forecasters(get_forecasters_request)
45+
get_forecasters_request = messages_pb2.ListForecastersRequest()
46+
get_forecasters_response = await client.ListForecasters(get_forecasters_request)
4647
forecasters = get_forecasters_response.forecasters
4748
return forecasters
4849

4950

5051
@dataclasses.dataclass
5152
class PageConfig:
52-
location: dp.ListLocationsResponseLocationSummary
53-
forecasters: list[dp.Forecaster]
53+
location: messages_pb2.ListLocationsResponse.LocationSummary
54+
forecasters: list[messages_pb2.Forecaster]
5455
start_date: dt.datetime
5556
end_date: dt.datetime
5657
forecast_type: str
@@ -62,7 +63,7 @@ class PageConfig:
6263
strict_horizon_filtering: bool
6364

6465

65-
async def setup_page(client: dp.DataPlatformDataServiceStub) -> PageConfig:
66+
async def setup_page(client: service_pb2_grpc.DataPlatformDataServiceStub) -> PageConfig:
6667
"""Setup the Streamlit page with sidebar options."""
6768
location_names = await get_location_names(client)
6869
selected_location_name = st.sidebar.selectbox(

‎src/dataplatform/toolbox/location.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,21 +11,21 @@ async def locations_section(data_client):
1111

1212
# Energy source and location type mappings
1313
ENERGY_SOURCES = {
14-
"UNSPECIFIED": dp.EnergySource.UNSPECIFIED,
15-
"SOLAR": dp.EnergySource.SOLAR,
16-
"WIND": dp.EnergySource.WIND,
14+
"UNSPECIFIED": common_pb2.EnergySource.ENERGY_SOURCE_UNSPECIFIED,
15+
"SOLAR": common_pb2.EnergySource.ENERGY_SOURCE_SOLAR,
16+
"WIND": common_pb2.EnergySource.ENERGY_SOURCE_WIND,
1717
}
1818

1919
LOCATION_TYPES = {
20-
"UNSPECIFIED": dp.LocationType.UNSPECIFIED,
21-
"SITE": dp.LocationType.SITE,
22-
"GSP": dp.LocationType.GSP,
23-
"DNO": dp.LocationType.DNO,
24-
"NATION": dp.LocationType.NATION,
25-
"STATE": dp.LocationType.STATE,
26-
"COUNTY": dp.LocationType.COUNTY,
27-
"CITY": dp.LocationType.CITY,
28-
"PRIMARY SUBSTATION": dp.LocationType.PRIMARY_SUBSTATION,
20+
"UNSPECIFIED": common_pb2.LocationType.LOCATION_TYPE_UNSPECIFIED,
21+
"SITE": common_pb2.LocationType.LOCATION_TYPE_SITE,
22+
"GSP": common_pb2.LocationType.LOCATION_TYPE_GSP,
23+
"DNO": common_pb2.LocationType.LOCATION_TYPE_DNO,
24+
"NATION": common_pb2.LocationType.LOCATION_TYPE_NATION,
25+
"STATE": common_pb2.LocationType.LOCATION_TYPE_STATE,
26+
"COUNTY": common_pb2.LocationType.LOCATION_TYPE_COUNTY,
27+
"CITY": common_pb2.LocationType.LOCATION_TYPE_CITY,
28+
"PRIMARY SUBSTATION": common_pb2.LocationType.LOCATION_TYPE_PRIMARY_SUBSTATION,
2929
}
3030

3131
# List Locations
@@ -50,15 +50,15 @@ async def locations_section(data_client):
5050
st.error("❌ Could not connect to Data Platform")
5151
else:
5252
try:
53-
request = dp.ListLocationsRequest()
53+
request = messages_pb2.ListLocationsRequest()
5454
if energy_source_filter != "UNSPECIFIED":
5555
request.energy_source_filter = ENERGY_SOURCES[energy_source_filter]
5656
if location_type_filter != "UNSPECIFIED":
5757
request.location_type_filter = LOCATION_TYPES[location_type_filter]
5858
if user_filter:
5959
request.user_oauth_id_filter = user_filter
6060

61-
response = await data_client.list_locations(request)
61+
response = await data_client.ListLocations(request)
6262
locations = response.locations
6363

6464
if locations:
@@ -92,8 +92,8 @@ async def locations_section(data_client):
9292
st.error("❌ Could not connect to Data Platform")
9393
else:
9494
try:
95-
response = await data_client.get_location(
96-
dp.GetLocationRequest(
95+
response = await data_client.GetLocation(
96+
messages_pb2.GetLocationRequest(
9797
location_uuid=loc_uuid,
9898
energy_source=ENERGY_SOURCES[loc_energy],
9999
include_geometry=include_geometry,
@@ -154,8 +154,8 @@ async def locations_section(data_client):
154154
try:
155155
# Parse metadata JSON
156156
metadata = json.loads(loc_metadata) if loc_metadata.strip() else {}
157-
response = await data_client.create_location(
158-
dp.CreateLocationRequest(
157+
response = await data_client.CreateLocation(
158+
messages_pb2.CreateLocationRequest(
159159
location_name=loc_name,
160160
energy_source=ENERGY_SOURCES.get(loc_energy_src, 1),
161161
location_type=LOCATION_TYPES.get(loc_type, 1),

0 commit comments

Comments
 (0)