changed logic to be cleaner; always create backups

Author: aysim319

nhsn/delphi_nhsn/pull.py

@@ -9,77 +9,56 @@
 from urllib.error import HTTPError
 
 import pandas as pd
-from delphi_epidata import Epidata
 from delphi_utils import create_backup_csv
-from epiweeks import Week
 from sodapy import Socrata
 
 from .constants import MAIN_DATASET_ID, PRELIM_DATASET_ID, PRELIM_SIGNALS_MAP, PRELIM_TYPE_DICT, SIGNALS_MAP, TYPE_DICT
 
 
-def get_latest_nhsn_issue_date() -> datetime:
-    """Check the api and return approxiate latest issue date available."""
-    response = Epidata.covidcast_meta()
-    meta_df = pd.DataFrame(response["epidata"])
-    nhsn_df = meta_df[meta_df["data_source"] == "nhsn"]
-    last_updated = datetime.utcfromtimestamp(min(nhsn_df.last_update))
-    max_issue = Week.fromstring(str(min(nhsn_df.max_issue)))
-    # last_updated can be multiple days off from max issue from patching / other processes
-    if max_issue.startdate() <= last_updated.date() <= max_issue.enddate():
-        return last_updated
-    # defaults to start of epiweek if last_update is vastly different from max_issue
-    return datetime.combine(max_issue.startdate(), datetime.min.time())
-
-
-def check_last_updated(client, dataset_id, logger):
+def check_last_updated(socrata_token, dataset_id, logger):
     """Check last updated timestamp to determine data should be pulled or not."""
-    # retry logic for 500 error
-    try:
-        response = client.get_metadata(dataset_id)
-    except HTTPError as err:
-        if err.code == 503:
-            time.sleep(2 + random.randint(0, 1000) / 1000.0)
-            response = client.get_metadata(dataset_id)
-        else:
-            raise err
+    client = Socrata("data.cdc.gov", socrata_token)
+    response = client.get_metadata(dataset_id)
 
     updated_timestamp = datetime.utcfromtimestamp(int(response["rowsUpdatedAt"]))
     now = datetime.utcnow()
     recently_updated_source = (now - updated_timestamp) < timedelta(days=1)
 
-    latest_issue_date = get_latest_nhsn_issue_date()
-    # generally expect one issue per week
-    recently_updated_api = (updated_timestamp - latest_issue_date) < timedelta(days=6)
-
     prelim_prefix = "Preliminary " if dataset_id == PRELIM_DATASET_ID else ""
     if recently_updated_source:
         logger.info(f"{prelim_prefix}NHSN data was recently updated; Pulling data", updated_timestamp=updated_timestamp)
-    elif not recently_updated_api:
-        logger.info(
-            f"{prelim_prefix}NHSN data is missing issue; Pulling data",
-            updated_timestamp=updated_timestamp,
-            issue=latest_issue_date,
-        )
     else:
         logger.info(f"{prelim_prefix}NHSN data is stale; Skipping", updated_timestamp=updated_timestamp)
-    return recently_updated_source or not recently_updated_api
+    return recently_updated_source
 
 
-def pull_data(socrata_token: str, dataset_id: str, logger):
+def pull_data(socrata_token: str, dataset_id: str, backup_dir:str, logger):
     """Pull data from Socrata API."""
     client = Socrata("data.cdc.gov", socrata_token)
-    recently_updated = check_last_updated(client, dataset_id, logger)
-    df = pd.DataFrame()
-    if recently_updated:
-        results = []
-        offset = 0
-        limit = 50000  # maximum limit allowed by SODA 2.0
+    logger.info("Pulling data from Socrata API")
+    results = []
+    offset = 0
+    limit = 50000  # maximum limit allowed by SODA 2.0
+    # retry logic for 500 error
+    try:
         page = client.get(dataset_id, limit=limit, offset=offset)
-        while page:
-            results.extend(page)
-            offset += limit
+    except HTTPError as err:
+        if err.code == 503:
+            time.sleep(2 + random.randint(0, 1000) / 1000.0)
             page = client.get(dataset_id, limit=limit, offset=offset)
+        else:
+            raise err
+
+    while len(page) > 0:
+        results.extend(page)
+        offset += limit
+        page = client.get(dataset_id, limit=limit, offset=offset)
+
+    if results:
         df = pd.DataFrame.from_records(results)
+        create_backup_csv(df, backup_dir, False, logger=logger)
+    else:
+        df = pd.DataFrame()
     return df
 
 
@@ -144,16 +123,16 @@ def pull_nhsn_data(
     """
     # Pull data from Socrata API
     df = (
-        pull_data(socrata_token, MAIN_DATASET_ID, logger)
+        pull_data(socrata_token, MAIN_DATASET_ID, backup_dir, logger)
         if not custom_run
         else pull_data_from_file(backup_dir, issue_date, logger, prelim_flag=False)
     )
 
-    keep_columns = list(TYPE_DICT.keys())
+    recently_updated = True if custom_run else check_last_updated(socrata_token, MAIN_DATASET_ID, logger)
 
-    if not df.empty:
-        create_backup_csv(df, backup_dir, custom_run, logger=logger)
+    keep_columns = list(TYPE_DICT.keys())
 
+    if not df.empty and recently_updated:
         df = df.rename(columns={"weekendingdate": "timestamp", "jurisdiction": "geo_id"})
 
         for signal, col_name in SIGNALS_MAP.items():
@@ -209,16 +188,15 @@ def pull_preliminary_nhsn_data(
     """
     # Pull data from Socrata API
     df = (
-        pull_data(socrata_token, PRELIM_DATASET_ID, logger)
+        pull_data(socrata_token, PRELIM_DATASET_ID, backup_dir, logger)
         if not custom_run
         else pull_data_from_file(backup_dir, issue_date, logger, prelim_flag=True)
     )
 
     keep_columns = list(PRELIM_TYPE_DICT.keys())
+    recently_updated = True if custom_run else check_last_updated(socrata_token, PRELIM_DATASET_ID, logger)
 
-    if not df.empty:
-        create_backup_csv(df, backup_dir, custom_run, sensor="prelim", logger=logger)
-
+    if not df.empty and recently_updated:
         df = df.rename(columns={"weekendingdate": "timestamp", "jurisdiction": "geo_id"})
 
         for signal, col_name in PRELIM_SIGNALS_MAP.items():