nssp patching code

nssp/README.md

@@ -78,3 +78,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 `bigchunk-dev-02.delphi.cmu.edu` under `/common/source_backup/nssp`. This data is needed to create patches. Talk to your sysadmin for access.
+When your credentials to the server 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,153 @@
+"""
+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_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_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
+
+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", "user"]
+        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 patch():
+    """
+    Run the 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)
+
+    current_issue = start_issue
+    while current_issue <= end_issue:
+        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)
+        current_issue += timedelta(days=1)
+
+    if download_source:
+        rmtree(source_dir)
+
+
+if __name__ == "__main__":
+    patch()

nssp/delphi_nssp/pull.py

@@ -1,10 +1,15 @@
 # -*- 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
 
@@ -20,6 +25,70 @@
 )
 
 
+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 = "delphi-master-prod-01.delphi.cmu.edu"
+    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"])
+    primary_source_files = [f"{date.strftime('%Y%m%d')}.csv.gz" for date in dates]
+    secondary_source_files = [f"{date.strftime('%Y%m%d')}_secondary.csv.gz" for date in dates]
+    remote_source_files = primary_source_files + secondary_source_files
+
+    # 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."""
     warn = textwrap.dedent(
@@ -48,10 +117,12 @@ 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
     """
+
     client = Socrata("data.cdc.gov", socrata_token)
     results = []
     offset = 0
@@ -65,8 +136,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
 
@@ -80,9 +157,22 @@ 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"
+        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)
 
@@ -99,7 +189,11 @@ def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, logger
 
 
 def secondary_pull_nssp_data(
-    socrata_token: str, backup_dir: str, custom_run: bool, logger: Optional[logging.Logger] = None
+    socrata_token: str,
+    backup_dir: str,
+    custom_run: bool,
+    issue_date: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
 ):
     """Pull the latest NSSP ER visits secondary dataset.
 
@@ -119,9 +213,23 @@ def secondary_pull_nssp_data(
     pd.DataFrame
         Dataframe as described above.
     """
-    socrata_results = pull_with_socrata_api(socrata_token, "7mra-9cq9")
-    df_ervisits = pd.DataFrame.from_records(socrata_results)
-    create_backup_csv(df_ervisits, backup_dir, custom_run, sensor="secondary", logger=logger)
+    if not custom_run:
+        socrata_results = pull_with_socrata_api(socrata_token, "7mra-9cq9")
+        df_ervisits = pd.DataFrame.from_records(socrata_results)
+        create_backup_csv(df_ervisits, backup_dir, custom_run, sensor="secondary", logger=logger)
+        logger.info("Number of records grabbed", num_records=len(df_ervisits), source="secondary 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}_secondary.csv.gz"
+        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=SECONDARY_COLS_MAP)
 
     # geo_type is not provided in the dataset, so we infer it from the geo_value