diff --git a/cwms/api.py b/cwms/api.py index 0784592e..c500c6ce 100644 --- a/cwms/api.py +++ b/cwms/api.py @@ -35,6 +35,7 @@ from requests import Response, adapters from requests_toolbelt import sessions # type: ignore from requests_toolbelt.sessions import BaseUrlSession # type: ignore +from urllib3.util.retry import Retry from cwms.cwms_types import JSON, RequestParams @@ -43,8 +44,24 @@ API_VERSION = 2 # Initialize a non-authenticated session with the default root URL and set default pool connections. + +retry_strategy = Retry( + total=6, + backoff_factor=0.5, + status_forcelist=[ + 403, + 429, + 500, + 502, + 503, + 504, + ], # Example: also retry on these HTTP status codes + allowed_methods=["GET", "PUT", "POST", "PATCH", "DELETE"], # Methods to retry +) SESSION = sessions.BaseUrlSession(base_url=API_ROOT) -adapter = adapters.HTTPAdapter(pool_connections=100, pool_maxsize=100) +adapter = adapters.HTTPAdapter( + pool_connections=100, pool_maxsize=100, max_retries=retry_strategy +) SESSION.mount("https://", adapter) @@ -119,7 +136,9 @@ def init_session( logging.debug(f"Initializing root URL: api_root={api_root}") SESSION = sessions.BaseUrlSession(base_url=api_root) adapter = adapters.HTTPAdapter( - pool_connections=pool_connections, pool_maxsize=pool_connections + pool_connections=pool_connections, + pool_maxsize=pool_connections, + max_retries=retry_strategy, ) SESSION.mount("https://", adapter) if api_key: diff --git a/cwms/timeseries/timeseries.py b/cwms/timeseries/timeseries.py index 1e72e5f3..82f8b52c 100644 --- a/cwms/timeseries/timeseries.py +++ b/cwms/timeseries/timeseries.py @@ -1,4 +1,4 @@ -import threading +import concurrent.futures from datetime import datetime from typing import Any, Dict, Optional @@ -16,6 +16,7 @@ def get_multi_timeseries_df( begin: Optional[datetime] = None, end: Optional[datetime] = None, melted: Optional[bool] = False, + max_workers: Optional[int] = 30, ) -> DataFrame: """gets multiple timeseries and stores into a single dataframe @@ -46,6 +47,9 @@ def get_multi_timeseries_df( melted: Boolean, optional, default is false if set to True a melted dataframe will be provided. By default a multi-index column dataframe will be returned. + max_workers: Int, Optional, default is None + It is a number of Threads aka size of pool in concurrent.futures.ThreadPoolExecutor. From 3.8 onwards + default value is min(32, os.cpu_count() + 4). Out of these 5 threads are preserved for I/O bound task. Returns @@ -53,60 +57,48 @@ def get_multi_timeseries_df( dataframe """ - def get_ts_ids( - result_dict: list[Dict[str, Any]], - ts_id: str, - office_id: str, - begin: datetime, - end: datetime, - unit: str, - version_date: datetime, - ) -> None: - data = get_timeseries( - ts_id=ts_id, - office_id=office_id, - unit=unit, - begin=begin, - end=end, - version_date=version_date, - ) - result_dict.append( - { - "ts_id": ts_id, - "unit": data.json["units"], - "version_date": version_date, - "values": data.df, - } - ) + def get_ts_ids(ts_id: str) -> Any: - result_dict = [] # type: list[Dict[str,Any]] - threads = [] - for ts_id in ts_ids: if ":" in ts_id: ts_id, version_date = ts_id.split(":", 1) version_date_dt = pd.to_datetime(version_date) else: version_date_dt = None - t = threading.Thread( - target=get_ts_ids, - args=(result_dict, ts_id, office_id, begin, end, unit, version_date_dt), - ) - threads.append(t) - t.start() + try: + data = get_timeseries( + ts_id=ts_id, + office_id=office_id, + unit=unit, + begin=begin, + end=end, + version_date=version_date_dt, + ) + result_dict = { + "ts_id": ts_id, + "unit": data.json["units"], + "version_date": version_date_dt, + "values": data.df, + } + return result_dict + except Exception as e: + print(f"Error processing {ts_id}: {e}") + return None - for t in threads: - t.join() + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + results = executor.map(get_ts_ids, ts_ids) + result_dict = list(results) data = pd.DataFrame() for row in result_dict: - temp_df = row["values"] - temp_df = temp_df.assign(ts_id=row["ts_id"], units=row["unit"]) - if "version_date" in row.keys(): - temp_df = temp_df.assign(version_date=row["version_date"]) - temp_df.dropna(how="all", axis=1, inplace=True) - data = pd.concat([data, temp_df], ignore_index=True) - - if not melted: + if row: + temp_df = row["values"] + temp_df = temp_df.assign(ts_id=row["ts_id"], units=row["unit"]) + if "version_date" in row.keys(): + temp_df = temp_df.assign(version_date=row["version_date"]) + temp_df.dropna(how="all", axis=1, inplace=True) + data = pd.concat([data, temp_df], ignore_index=True) + + if not melted and "date-time" in data.columns: cols = ["ts_id", "units"] if "version_date" in data.columns: cols.append("version_date") @@ -129,7 +121,7 @@ def get_timeseries( datum: Optional[str] = None, begin: Optional[datetime] = None, end: Optional[datetime] = None, - page_size: Optional[int] = 500000, + page_size: Optional[int] = 300000, version_date: Optional[datetime] = None, trim: Optional[bool] = True, ) -> Data: @@ -163,7 +155,7 @@ def get_timeseries( not specified, any required time window ends at the current time. Any timezone information should be passed within the datetime object. If no timezone information is given, default will be UTC. - page_size: int, optional, default is 5000000: Specifies the number of records to obtain in + page_size: int, optional, default is 300000: Specifies the number of records to obtain in a single call. version_date: datetime, optional, default is None Version date of time series values being requested. If this field is not specified and @@ -242,7 +234,7 @@ def timeseries_df_to_json( df["quality-code"] = 0 if "date-time" not in df: raise TypeError( - "date-time is a required column in data when posting as a dateframe" + "date-time is a required column in data when posting as a dataframe" ) if "value" not in df: raise TypeError( diff --git a/tests/timeseries/timeseries_test.py b/tests/timeseries/timeseries_test.py index 6672c821..03e64c91 100644 --- a/tests/timeseries/timeseries_test.py +++ b/tests/timeseries/timeseries_test.py @@ -148,7 +148,7 @@ def test_get_timeseries_unversioned_default(requests_mock): "unit=EN&" "begin=2008-05-01T15%3A00%3A00%2B00%3A00&" "end=2008-05-01T17%3A00%3A00%2B00%3A00&" - "page-size=500000", + "page-size=300000", json=_UNVERS_TS_JSON, ) @@ -177,7 +177,7 @@ def test_get_empty_ts_df(requests_mock): "unit=EN&" "begin=2008-05-01T15%3A00%3A00%2B00%3A00&" "end=2008-05-01T17%3A00%3A00%2B00%3A00&" - "page-size=500000&" + "page-size=300000&" "trim=true", json=_EMPTY_TS_JSON, ) @@ -306,7 +306,7 @@ def test_get_multi_timeseries_default(requests_mock): "unit=EN&" "begin=2008-05-01T15%3A00%3A00%2B00%3A00&" "end=2008-05-01T17%3A00%3A00%2B00%3A00&" - "page-size=500000&" + "page-size=300000&" "version-date=2021-06-20T08%3A00%3A00%2B00%3A00", json=_VERS_TS_JSON, ) @@ -318,7 +318,7 @@ def test_get_multi_timeseries_default(requests_mock): "unit=EN&" "begin=2008-05-01T15%3A00%3A00%2B00%3A00&" "end=2008-05-01T17%3A00%3A00%2B00%3A00&" - "page-size=500000", + "page-size=300000", json=_UNVERS_TS_JSON, ) @@ -362,7 +362,7 @@ def test_get_timeseries_versioned_default(requests_mock): "unit=EN&" "begin=2008-05-01T15%3A00%3A00%2B00%3A00&" "end=2008-05-01T17%3A00%3A00%2B00%3A00&" - "page-size=500000&" + "page-size=300000&" "version-date=2021-06-20T08%3A00%3A00%2B00%3A00", json=_VERS_TS_JSON, )