fix(partial records): don't persist partial records

Author: jacob-a-brown

backend/connectors/usgs/source.py

@@ -32,7 +32,7 @@
     NWISSiteTransformer,
     NWISWaterLevelTransformer,
 )
-from backend.exceptions import USGSRateLimitError
+from backend.exceptions import USGSRateLimitError, PartialOrNoDataError
 from backend.source import (
     BaseWaterLevelSource,
     BaseSiteSource,
@@ -91,6 +91,7 @@ def get_records(self):
         # if config.end_date:
         #     params["endDt"] = config.end_dt.date().isoformat()
 
+        data: dict = {}
         tries: int = 0
 
         while tries < MAX_RETRIES:
@@ -107,7 +108,7 @@ def get_records(self):
                 )
 
                 if response.status_code == 200:
-                    data: dict = response.json()
+                    data = response.json()
                     break
                 elif response.status_code == 429:
                     raise USGSRateLimitError()
@@ -125,6 +126,10 @@ def get_records(self):
             tries += 1
             time.sleep(tries)
 
+        if data == {}:
+            self.warn("Failed to retrieve site records after multiple attempts.")
+            raise PartialOrNoDataError("Failed to retrieve site records after multiple attempts.")
+
         records: list = data.get("features", [])
 
         return records
@@ -168,6 +173,8 @@ def get_records(self, site_record):
                     list_of_sites
                 ]
             }
+
+            data: dict = {}
             tries: int = 0
 
             while tries < MAX_RETRIES:
@@ -184,7 +191,7 @@ def get_records(self, site_record):
                         timeout=TIMEOUT,
                     )
                     if response.status_code == 200:
-                        data: dict = response.json()
+                        data = response.json()
                         break
                     elif response.status_code == 429:
                         raise USGSRateLimitError()
@@ -202,6 +209,10 @@ def get_records(self, site_record):
                 tries += 1
                 time.sleep(tries)
 
+            if data == {}:
+                self.warn("Failed to retrieve water level records after multiple attempts.")
+                raise PartialOrNoDataError("Failed to retrieve water level records after multiple attempts.")
+
             features: list[dict] = data.get("features", [])
 
             standard_features: list[dict] = [self._standardize_record(feature) for feature in features]
@@ -225,6 +236,7 @@ def get_records(self, site_record):
             # USGS APIs use cursor pagination, so we can just follow the "next" links until there are no more
             while found_next_link:
                 tries: int = 0
+                data: dict = {}
                 while tries < MAX_RETRIES:
                     try:
                         response = httpx.get(
@@ -233,7 +245,7 @@ def get_records(self, site_record):
                                 timeout=TIMEOUT,
                             )
                         if response.status_code == 200:
-                            data: dict = response.json()
+                            data = response.json()
                             break
                         elif response.status_code == 429:
                             raise USGSRateLimitError()
@@ -250,6 +262,10 @@ def get_records(self, site_record):
                     tries += 1
                     time.sleep(tries)
                 
+                if data == {}:
+                    self.warn("Failed to retrieve paginated water level records after multiple attempts.")
+                    raise PartialOrNoDataError("Failed to retrieve paginated water level records after multiple attempts
+
                 features: list[dict] = data.get("features", [])
                 standard_features: list[dict] = [self._standardize_record(feature) for feature in features]
                 records.extend(standard_features)

backend/exceptions.py

@@ -1,2 +1,5 @@
 class USGSRateLimitError(Exception):
+    pass
+
+class PartialOrNoDataError(Exception):
     pass

backend/unifier.py

@@ -21,7 +21,7 @@
 from backend.persister import BasePersister
 from backend.persisters.geoserver import GeoServerPersister
 from backend.source import BaseSiteSource
-from backend.exceptions import USGSRateLimitError
+from backend.exceptions import USGSRateLimitError, PartialOrNoDataError
 
 
 def health_check(source: BaseSiteSource) -> bool | None:
@@ -171,8 +171,8 @@ def _site_wrapper(site_source, parameter_source, persister, config):
                         summary_records = parameter_source.read(
                             site_records, use_summarize, start_ind, end_ind
                         )
-                    except USGSRateLimitError:
-                        # if a rate limit error is hit we want to remove USGS sites so there aren't partial records
+                    except (USGSRateLimitError, PartialOrNoDataError):
+                        # remove partial records to prevent incomplete data from being saved
                         persister.sites = persister.sites[:initial_sites_len]
                         persister.timeseries = persister.timeseries[:initial_timeseries_len]
                         persister.records = persister.records[:initial_records_len]
@@ -187,8 +187,8 @@ def _site_wrapper(site_source, parameter_source, persister, config):
                         results = parameter_source.read(
                             site_records, use_summarize, start_ind, end_ind
                         )
-                    except USGSRateLimitError:
-                        # if a rate limit error is hit we want to remove USGS sites so there aren't partial records
+                    except (USGSRateLimitError, PartialOrNoDataError):
+                        # remove partial records to prevent incomplete data from being saved
                         persister.sites = persister.sites[:initial_sites_len]
                         persister.timeseries = persister.timeseries[:initial_timeseries_len]
                         break