Merge pull request #2000 from cmu-delphi/nssp_patching

nssp patching code

Author: minhkhul

_delphi_utils_python/delphi_utils/export.py

@@ -205,4 +205,5 @@ def create_backup_csv(
                 )
         # pylint: disable=W0703
         except Exception as e:
-            logger.info("Backup file creation failed", msg=e)
+            if logger:
+                logger.info("Backup file creation failed", msg=e)

nssp/README.md

@@ -76,3 +76,10 @@ with the percentage of code covered by the tests.
 
 None of the linting or unit tests should fail, and the code lines that are not covered by unit tests should be small and
 should not include critical sub-routines. 
+
+## Running Patches:
+A daily backup of source in the form of csv files can be found on prod under `/common/covidcast/source_backup/nssp/`. This data is needed to create patches. Talk to your sysadmin for access.
+When your credentials are working, to create patching data for a specific date range in batch issue format, adjust `params.json` in accordance with instructions in `patch.py`, then run
+```
+env/bin/python -m delphi_nssp.patch
+```

nssp/delphi_nssp/patch.py

@@ -0,0 +1,186 @@
+"""
+This module is used for patching data in the delphi_nssp package.
+
+The code assume user can use key-based auth to access prod server
+where historical source data is stored.
+
+To use this module, configure params.json like so:
+
+{
+  "common": {
+    "custom_run": true,
+    ...
+  },
+  "validation": {
+    ...
+  },
+  "patch": {
+    "source_host": "prod.server.edu",
+    "source_dir": "delphi/covidcast-indicators/nssp/source_data",
+    "user": "username",
+    "patch_dir": "delphi/covidcast-indicators/nssp/AprilPatch",
+    "start_issue": "2024-04-20",
+    "end_issue": "2024-04-21",
+  }
+}
+
+In this params.json, we
+- Turn on the "custom_run" flag under "common"
+- Add "patch" section, which contains:
+    + "source_host": the prod server where source data is backed up
+    + "source_dir": the local directory where source data is downloaded to
+    + "user": the username to log in to the remote server where source data is backed up
+    + "patch_dir": the local directory where to write all patch issues output
+    + "start_date": str, YYYY-MM-DD format, first issue date
+    + "end_date": str, YYYY-MM-DD format, last issue date
+
+if "source_dir" doesn't exist locally or has no files in it, we download source data to source_dir
+else, we assume all needed source files are already in source_dir.
+
+This module will generate data for that range of issue dates, and store them in batch issue format in the patch_dir:
+[patch_dir]/issue_[issue-date]/nssp/actual_data_file.csv
+"""
+
+import sys
+from datetime import datetime, timedelta
+from os import listdir, makedirs, path
+from shutil import rmtree
+
+import pandas as pd
+from delphi_utils import get_structured_logger, read_params
+from epiweeks import Week
+
+from .pull import get_source_data
+from .run import run_module
+
+
+def good_patch_config(params, logger):
+    """
+    Check if the params.json file is correctly configured for patching.
+
+    params: Dict[str, Any]
+        Nested dictionary of parameters, typically loaded from params.json file.
+    logger: Logger object
+        Logger object to log messages.
+    """
+    valid_config = True
+    custom_run = params["common"].get("custom_run", False)
+    if not custom_run:
+        logger.error("Calling patch.py without custom_run flag set true.")
+        valid_config = False
+
+    patch_config = params.get("patch", {})
+    if patch_config == {}:
+        logger.error("Custom flag is on, but patch section is missing.")
+        valid_config = False
+    else:
+        required_patch_keys = ["start_issue", "end_issue", "patch_dir", "source_dir"]
+
+        source_dir = params["patch"]["source_dir"]
+        if not path.isdir(source_dir) or not listdir(source_dir):
+            required_patch_keys.append("user")
+            required_patch_keys.append("source_host")
+
+        missing_keys = [key for key in required_patch_keys if key not in patch_config]
+        if missing_keys:
+            logger.error("Patch section is missing required key(s)", missing_keys=missing_keys)
+            valid_config = False
+        else:
+            try:  # issue dates validity check
+                start_issue = datetime.strptime(patch_config["start_issue"], "%Y-%m-%d")
+                end_issue = datetime.strptime(patch_config["end_issue"], "%Y-%m-%d")
+                if start_issue > end_issue:
+                    logger.error("Start issue date is after end issue date.")
+                    valid_config = False
+            except ValueError:
+                logger.error("Issue dates must be in YYYY-MM-DD format.")
+                valid_config = False
+
+    if valid_config:
+        logger.info("Good patch configuration.")
+        return True
+    logger.info("Bad patch configuration.")
+    return False
+
+
+def get_patch_dates(start_issue, end_issue, source_dir):
+    """
+    Get the dates to run patch on given a range of issue dates.
+
+    Due to weekly cadence of nssp data, dates to run patch on are not necessarily the same as issue dates.
+    We use the latest date with source data per epiweek as reporting date for patching of that week's data.
+
+    start_issue: datetime object
+    end_issue: datetime object
+    """
+    patch_dates = []
+    date_range = pd.date_range(start=start_issue, end=end_issue)
+    dates_with_source_data = {
+        date for date in date_range if path.isfile(f"""{source_dir}/{date.strftime("%Y%m%d")}.csv.gz""")
+    }
+    epiweek_start_dates = {Week.fromdate(date).startdate() for date in date_range}
+    for epiweek_start_date in epiweek_start_dates:
+        epiweek = Week.fromdate(epiweek_start_date)
+        dates_with_data_in_epiweek = [date for date in dates_with_source_data if date.date() in epiweek.iterdates()]
+        if dates_with_data_in_epiweek == []:
+            continue
+        latest_date_with_data = max(dates_with_data_in_epiweek)
+        patch_dates.append(latest_date_with_data)
+    patch_dates.sort()
+    return patch_dates
+
+
+def patch():
+    """Run nssp indicator for a range of issue dates."""
+    params = read_params()
+    logger = get_structured_logger("delphi_nssp.patch", filename=params["common"]["log_filename"])
+    if not good_patch_config(params, logger):
+        sys.exit(1)
+
+    source_dir = params["patch"]["source_dir"]
+    download_source = False
+    if not path.isdir(source_dir) or not listdir(source_dir):  # no source dir or empty source dir
+        download_source = True
+        get_source_data(params, logger)
+    else:
+        logger.info("Source data already exists locally.")
+
+    start_issue = datetime.strptime(params["patch"]["start_issue"], "%Y-%m-%d")
+    end_issue = datetime.strptime(params["patch"]["end_issue"], "%Y-%m-%d")
+
+    logger.info(start_issue=start_issue.strftime("%Y-%m-%d"))
+    logger.info(end_issue=end_issue.strftime("%Y-%m-%d"))
+    logger.info(source_dir=source_dir)
+    logger.info(patch_dir=params["patch"]["patch_dir"])
+    makedirs(params["patch"]["patch_dir"], exist_ok=True)
+
+    patch_dates = get_patch_dates(start_issue, end_issue, source_dir)
+
+    for current_issue in patch_dates:
+        logger.info("patching issue", issue_date=current_issue.strftime("%Y%m%d"))
+
+        current_issue_source_csv = f"""{source_dir}/{current_issue.strftime("%Y%m%d")}.csv.gz"""
+        if not path.isfile(current_issue_source_csv):
+            logger.info("No source data at this path", current_issue_source_csv=current_issue_source_csv)
+            current_issue += timedelta(days=1)
+            continue
+
+        params["patch"]["current_issue"] = current_issue.strftime("%Y%m%d")
+
+        # current_issue_date can be different from params["patch"]["current_issue"]
+        # due to weekly cadence of nssp data. For weekly sources, issue dates in our
+        # db matches with first date of epiweek that the reporting date falls in,
+        # rather than reporting date itself.
+        current_issue_date = Week.fromdate(current_issue).startdate()
+        current_issue_dir = f"""{params["patch"]["patch_dir"]}/issue_{current_issue_date.strftime("%Y%m%d")}/nssp"""
+        makedirs(f"{current_issue_dir}", exist_ok=True)
+        params["common"]["export_dir"] = f"""{current_issue_dir}"""
+
+        run_module(params, logger)
+
+    if download_source:
+        rmtree(source_dir)
+
+
+if __name__ == "__main__":
+    patch()

nssp/delphi_nssp/pull.py

@@ -1,20 +1,84 @@
 # -*- coding: utf-8 -*-
 """Functions for pulling NSSP ER data."""
+
+import functools
 import logging
+import sys
 import textwrap
+from os import makedirs, path
 from typing import Optional
 
 import pandas as pd
+import paramiko
 from delphi_utils import create_backup_csv
 from sodapy import Socrata
 
-from .constants import (
-    NEWLINE,
-    SIGNALS,
-    SIGNALS_MAP,
-    TYPE_DICT,
-)
+from .constants import NEWLINE, SIGNALS, SIGNALS_MAP, TYPE_DICT
+
+
+def print_callback(remote_file_name, logger, bytes_so_far, bytes_total, progress_chunks):
+    """Print the callback information."""
+    rough_percent_transferred = int(100 * (bytes_so_far / bytes_total))
+    if rough_percent_transferred in progress_chunks:
+        logger.info("Transfer in progress", remote_file_name=remote_file_name, percent=rough_percent_transferred)
+        # Remove progress chunk, so it is not logged again
+        progress_chunks.remove(rough_percent_transferred)
+
+
+def get_source_data(params, logger):
+    """
+    Download historical source data from a backup server.
+
+    This function is typically used in patching only. Normal runs grab latest data from SODA API.
+
+    This function uses "user" configuration under "patch" section in params.json to specify
+    a username with local key-based access to connect to server where backup nssp source data is stored.
+    It uses "backup_dir" config under "common" section to locate backup files on remote server.
+    It then searches for CSV files that match the inclusive range of issue dates
+    specified by 'start_issue', and 'end_issue' config.
 
+    These CSV files are then downloaded and stored in the local 'source_dir' directory.
+    """
+    makedirs(params["patch"]["source_dir"], exist_ok=True)
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    host = params["patch"]["source_host"]
+    user = params["patch"]["user"]
+    ssh.connect(host, username=user)
+
+    # Generate file names of source files to download
+    dates = pd.date_range(start=params["patch"]["start_issue"], end=params["patch"]["end_issue"])
+    remote_source_files = [f"{date.strftime('%Y%m%d')}.csv.gz" for date in dates]
+
+    # Download source files
+    sftp = ssh.open_sftp()
+    try:
+        sftp.stat(params["common"]["backup_dir"])
+    except IOError:
+        logger.error("Source backup directory does not exist on the remote server.")
+
+    sftp.chdir(params["common"]["backup_dir"])
+
+    num_files_transferred = 0
+    for remote_file_name in remote_source_files:
+        callback_for_filename = functools.partial(print_callback, remote_file_name, logger, progress_chunks=[0, 50])
+        local_file_path = path.join(params["patch"]["source_dir"], remote_file_name)
+        try:
+            sftp.stat(remote_file_name)
+        except IOError:
+            logger.warning(
+                "Source backup for this date does not exist on the remote server.",
+                missing_filename=remote_file_name,
+            )
+            continue
+        sftp.get(remote_file_name, local_file_path, callback=callback_for_filename)
+        logger.info("Transfer finished", remote_file_name=remote_file_name, local_file_path=local_file_path)
+        num_files_transferred += 1
+    ssh.close()
+
+    if num_files_transferred == 0:
+        logger.error("No source data was transferred. Check the source backup server for potential issues.")
+        sys.exit(1)
 
 def warn_string(df, type_dict):
     """Format the warning string."""
@@ -44,6 +108,7 @@ def pull_with_socrata_api(socrata_token: str, dataset_id: str):
     dataset_id: str
         The dataset id to pull data from
 
+
     Returns
     -------
     list of dictionaries, each representing a row in the dataset
@@ -61,8 +126,14 @@ def pull_with_socrata_api(socrata_token: str, dataset_id: str):
     return results
 
 
-def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, logger: Optional[logging.Logger] = None):
-    """Pull the latest NSSP ER visits primary dataset.
+def pull_nssp_data(
+    socrata_token: str,
+    backup_dir: str,
+    custom_run: bool,
+    issue_date: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+):
+    """Pull the NSSP ER visits primary dataset.
 
     https://data.cdc.gov/Public-Health-Surveillance/NSSP-Emergency-Department-Visit-Trajectories-by-St/rdmq-nq56/data_preview
 
@@ -76,9 +147,25 @@ def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, logger
     pd.DataFrame
         Dataframe as described above.
     """
-    socrata_results = pull_with_socrata_api(socrata_token, "rdmq-nq56")
-    df_ervisits = pd.DataFrame.from_records(socrata_results)
-    create_backup_csv(df_ervisits, backup_dir, custom_run, logger=logger)
+    if not custom_run:
+        socrata_results = pull_with_socrata_api(socrata_token, "rdmq-nq56")
+        df_ervisits = pd.DataFrame.from_records(socrata_results)
+        create_backup_csv(df_ervisits, backup_dir, custom_run, logger=logger)
+        logger.info("Number of records grabbed", num_records=len(df_ervisits), source="Socrata API")
+    elif custom_run and logger.name == "delphi_nssp.patch":
+        if issue_date is None:
+            raise ValueError("Issue date is required for patching")
+        source_filename = f"{backup_dir}/{issue_date}.csv.gz"
+        if not path.isfile(source_filename):
+            logger.warning("No primary source data found", source=source_filename, issue_date=issue_date)
+            return None
+        df_ervisits = pd.read_csv(source_filename)
+        logger.info(
+            "Number of records grabbed",
+            num_records=len(df_ervisits),
+            source=source_filename,
+        )
+
     df_ervisits = df_ervisits.rename(columns={"week_end": "timestamp"})
     df_ervisits = df_ervisits.rename(columns=SIGNALS_MAP)