feat: adding patching with backfill

Author: aysim319

claims_hosp/delphi_claims_hosp/backfill.py

@@ -5,15 +5,18 @@
 Created: 2022-08-03
 
 """
-import os
+
 import glob
+import os
+import re
+import shutil
 from datetime import datetime
+from typing import Union
 
 # third party
 import pandas as pd
 from delphi_utils import GeoMapper
 
-
 from .config import Config
 
 gmpr = GeoMapper()
@@ -69,9 +72,58 @@ def store_backfill_file(claims_filepath, _end_date, backfill_dir):
         "/claims_hosp_as_of_%s.parquet"%datetime.strftime(_end_date, "%Y%m%d")
     # Store intermediate file into the backfill folder
     backfilldata.to_parquet(path, index=False)
+    return path
+
+
+def merge_existing_backfill_files(backfill_dir, backfill_file, issue_date, logger):
+    """
+    Merge existing backfill with the patch data included.
+    Parameters
+    ----------
+    issue_date : datetime
+        The most recent date when the raw data is received
+    backfill_dir : str
+        specified path to store backfill files.
+    backfill_file : str
+
+    """
+
+    new_files = glob.glob(backfill_dir + "/claims_hosp_*")
+
+    def get_file_with_date(files) -> Union[str, None]:
+        for filename in files:
+            pattern = re.findall(r"\d{8}", filename)
+            if len(pattern) == 2:
+                start_date = datetime.strptime(pattern[0], "%Y%m%d")
+                end_date = datetime.strptime(pattern[1], "%Y%m%d")
+                if start_date <= issue_date or end_date <= issue_date:
+                    return filename
+        return ""
+
+    file_name = get_file_with_date(new_files)
+
+    if len(file_name) == 0:
+        logger.info("patch file is too recent to merge", issue_date=issue_date.strftime("%Y-%m-%d"))
+        return
+
+    # Start to merge files
+    merge_file = f"{file_name.split('.')[0]}_after_merge.parquet"
+    try:
+        shutil.copyfile(file_name, merge_file)
+        existing_df = pd.read_parquet(merge_file, engine="pyarrow")
+        df = pd.read_parquet(backfill_file, engine="pyarrow")
+        merged_df = pd.concat([existing_df, df]).sort_values(["time_value", "fips"])
+        merged_df.to_parquet(merge_file, index=False)
+        os.remove(file_name)
+        os.rename(merge_file, file_name)
+    # pylint: disable=W0703:
+    except Exception as e:
+        os.remove(merge_file)
+        logger.error(e)
+    return
+
 
-def merge_backfill_file(backfill_dir, backfill_merge_day, today,
-                        test_mode=False, check_nd=25):
+def merge_backfill_file(backfill_dir, backfill_merge_day, most_recent, test_mode=False, check_nd=25):
     """
     Merge ~4 weeks' backfill data into one file.
 
@@ -80,7 +132,7 @@ def merge_backfill_file(backfill_dir, backfill_merge_day, today,
     threshold to allow flexibility in data delivery.
     Parameters
     ----------
-    today : datetime
+    most_recent : datetime
         The most recent date when the raw data is received
     backfill_dir : str
         specified path to store backfill files.
@@ -109,7 +161,7 @@ def get_date(file_link):
 
     # Check whether to merge
     # Check the number of files that are not merged
-    if today.weekday() != backfill_merge_day or (today-earliest_date).days <= check_nd:
+    if most_recent.weekday() != backfill_merge_day or (most_recent - earliest_date).days <= check_nd:
         return
 
     # Start to merge files

claims_hosp/delphi_claims_hosp/run.py

@@ -5,22 +5,24 @@
 when the module is run with `python -m delphi_claims_hosp`.
 """
 
+import os
+
 # standard packages
 import time
-import os
 from datetime import datetime, timedelta
 from pathlib import Path
 
 # third party
 from delphi_utils import get_structured_logger
 
+from .backfill import merge_backfill_file, merge_existing_backfill_files, store_backfill_file
+
 # first party
 from .config import Config
 from .download_claims_ftp_files import download
-from .modify_claims_drops import modify_and_write
 from .get_latest_claims_name import get_latest_filename
+from .modify_claims_drops import modify_and_write
 from .update_indicator import ClaimsHospIndicatorUpdater
-from .backfill import (store_backfill_file, merge_backfill_file)
 
 
 def run_module(params, logger=None):
@@ -56,6 +58,10 @@ def run_module(params, logger=None):
     start_time = time.time()
     issue_date_str = params.get("patch", {}).get("current_issue", None)
     issue_date = datetime.strptime(issue_date_str, "%Y-%m-%d")
+    # safety check for patch parameters exists in file, but not running custom runs/patches
+    custom_run_flag = (
+        False if not params["indicator"].get("custom_run", False) else params["indicator"].get("custom_run", False)
+    )
     if not logger:
         logger = get_structured_logger(
             __name__,
@@ -64,8 +70,7 @@ def run_module(params, logger=None):
         )
 
     # pull latest data
-    download(params["indicator"]["ftp_credentials"],
-             params["indicator"]["input_dir"], logger, issue_date=issue_date)
+    download(params["indicator"]["ftp_credentials"], params["indicator"]["input_dir"], logger, issue_date=issue_date)
 
     # aggregate data
     modify_and_write(params["indicator"]["input_dir"], logger)
@@ -99,8 +104,13 @@ def run_module(params, logger=None):
     if params["indicator"].get("generate_backfill_files", True):
         backfill_dir = params["indicator"]["backfill_dir"]
         backfill_merge_day = params["indicator"]["backfill_merge_day"]
-        merge_backfill_file(backfill_dir, backfill_merge_day, datetime.today())
-        store_backfill_file(claims_file, dropdate_dt, backfill_dir)
+        if custom_run_flag:
+            backfilled_filepath = store_backfill_file(claims_file, dropdate_dt, backfill_dir)
+            merge_existing_backfill_files(backfill_dir, backfilled_filepath, issue_date, logger)
+
+        else:
+            merge_backfill_file(backfill_dir, backfill_merge_day, datetime.today())
+            store_backfill_file(claims_file, dropdate_dt, backfill_dir)
 
     # print out information
     logger.info("Loaded params",

claims_hosp/tests/test_backfill.py

@@ -1,27 +1,32 @@
+import logging
 import os
 import glob
 from datetime import datetime
+from pathlib import Path
+import shutil
 
 # third party
 import pandas as pd
 import pytest
 
 # first party
 from delphi_claims_hosp.config import Config, GeoConstants
-from delphi_claims_hosp.backfill import store_backfill_file, merge_backfill_file
+from delphi_claims_hosp.backfill import store_backfill_file, merge_backfill_file, merge_existing_backfill_files
 
 CONFIG = Config()
 CONSTANTS = GeoConstants()
+TEST_PATH = Path(__file__).parent
 PARAMS = {
     "indicator": {
-        "input_file": "test_data/SYNEDI_AGG_INPATIENT_11062020_1451CDT.csv.gz",
-        "backfill_dir": "./backfill",
+        "input_file": f"{TEST_PATH}/test_data/SYNEDI_AGG_INPATIENT_11062020_1451CDT.csv.gz",
+        "backfill_dir": f"{TEST_PATH}/backfill",
         "drop_date": "2020-06-11",
     }
 }
 DATA_FILEPATH = PARAMS["indicator"]["input_file"]
 DROP_DATE = pd.to_datetime(PARAMS["indicator"]["drop_date"])
 backfill_dir = PARAMS["indicator"]["backfill_dir"]
+TEST_LOGGER = logging.getLogger()
 
 class TestBackfill:
 
@@ -95,3 +100,67 @@ def test_merge_backfill_file(self):
 
         os.remove(backfill_dir + "/" + fn)
         assert fn not in os.listdir(backfill_dir)
+
+    def test_merge_existing_backfill_files(self):
+        issue_date = datetime(year=2020, month=6, day=13)
+        issue_date_str = issue_date.strftime("%Y%m%d")
+        def prep_backfill_data():
+            # Generate backfill daily files
+            for d in range(11, 15):
+                dropdate = datetime(2020, 6, d)
+                store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir)
+
+            today = datetime(2020, 6, 14)
+            # creating expected file
+            merge_backfill_file(backfill_dir, today.weekday(), today,
+                                test_mode=True, check_nd=2)
+            original = f"{backfill_dir}/claims_hosp_from_20200611_to_20200614.parquet"
+            os.rename(original, f"{backfill_dir}/expected.parquet")
+
+            # creating backfill without issue date
+            os.remove(f"{backfill_dir}/claims_hosp_as_of_{issue_date_str}.parquet")
+            today = datetime(2020, 6, 14)
+            merge_backfill_file(backfill_dir, today.weekday(), today,
+                                test_mode=True, check_nd=2)
+
+            old_files = glob.glob(backfill_dir + "/claims_hosp_as_of_*")
+            for file in old_files:
+                os.remove(file)
+
+        prep_backfill_data()
+        file_to_add = store_backfill_file(DATA_FILEPATH, issue_date, backfill_dir)
+        merge_existing_backfill_files(backfill_dir, file_to_add, issue_date, TEST_LOGGER)
+
+        expected = pd.read_parquet(f"{backfill_dir}/expected.parquet")
+        merged = pd.read_parquet(f"{backfill_dir}/claims_hosp_from_20200611_to_20200614.parquet")
+
+        check_diff = expected.merge(merged, how='left', indicator=True)
+        assert check_diff[check_diff["_merge"] == "both"].shape[0] == expected.shape[0]
+        for file in glob.glob(backfill_dir + "/*.parquet"):
+            os.remove(file)
+
+
+    def test_merge_existing_backfill_files_no_call(self):
+        issue_date = datetime(year=2020, month=6, day=20)
+        issue_date_str = issue_date.strftime("%Y%m%d")
+        def prep_backfill_data():
+            # Generate backfill daily files
+            for d in range(11, 15):
+                dropdate = datetime(2020, 6, d)
+                store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir)
+
+            today = datetime(2020, 6, 14)
+            # creating expected file
+            merge_backfill_file(backfill_dir, today.weekday(), today,
+                                test_mode=True, check_nd=8)
+
+        prep_backfill_data()
+        file_to_add = store_backfill_file(DATA_FILEPATH, issue_date, backfill_dir)
+        merge_existing_backfill_files(backfill_dir, file_to_add, issue_date, TEST_LOGGER)
+
+        old_files = glob.glob(backfill_dir + "*.parquet")
+        for file in old_files:
+            os.remove(file)
+
+
+

claims_hosp/tests/test_download_claims_ftp_files.py

@@ -1,19 +1,40 @@
 # standard
 import datetime
 import re
+from mock import MagicMock, patch
+import logging
 
 # third party
 import numpy as np
+from freezegun import freeze_time
 
 # first party
 from delphi_claims_hosp.download_claims_ftp_files import (change_date_format,
-                                                          get_timestamp)
+                                                          get_timestamp, download)
 
 OLD_FILENAME_TIMESTAMP = re.compile(
     r".*EDI_AGG_INPATIENT_[0-9]_(?P<ymd>[0-9]*)_(?P<hm>[0-9]*)[^0-9]*")
 NEW_FILENAME_TIMESTAMP = re.compile(r".*EDI_AGG_INPATIENT_(?P<ymd>[0-9]*)_(?P<hm>[0-9]*)[^0-9]*")
 
+TEST_LOGGER = logging.getLogger()
 class TestDownloadClaimsFtpFiles:
+
+    @patch('delphi_claims_hosp.download_claims_ftp_files.paramiko.SSHClient')
+    @patch('delphi_claims_hosp.download_claims_ftp_files.path.exists', return_value=False)
+    def test_download(self, mock_exists, mock_sshclient):
+        mock_sshclient_instance = MagicMock()
+        mock_sshclient.return_value = mock_sshclient_instance
+        mock_sftp = MagicMock()
+        mock_sshclient_instance.open_sftp.return_value = mock_sftp
+        mock_sftp.listdir_attr.return_value = [MagicMock(filename="SYNEDI_AGG_INPATIENT_11062020_1451CDT.csv.gz")]
+        ftp_credentials = {"host": "test_host", "user": "test_user", "pass": "test_pass", "port": "test_port"}
+        out_path = "./test_data/"
+
+        issue_date = datetime.datetime(2020, 11, 7)
+        download(ftp_credentials, out_path, TEST_LOGGER, issue_date=issue_date)
+        mock_sshclient_instance.connect.assert_called_once_with(ftp_credentials["host"], username=ftp_credentials["user"], password=ftp_credentials["pass"], port=ftp_credentials["port"])
+        mock_sftp.get.assert_called()
+
 
     def test_change_date_format(self):
         name = "SYNEDI_AGG_INPATIENT_20200611_1451CDT"

claims_hosp/tests/test_patch.py

@@ -0,0 +1,31 @@
+import unittest
+from unittest.mock import patch as mock_patch
+from delphi_claims_hosp.patch import patch
+import os
+import shutil
+
+class TestPatchModule:
+    def test_patch(self):
+        with mock_patch('delphi_claims_hosp.patch.get_structured_logger') as mock_get_structured_logger, \
+             mock_patch('delphi_claims_hosp.patch.read_params') as mock_read_params, \
+             mock_patch('delphi_claims_hosp.patch.run_module') as mock_run_module:
+
+            mock_read_params.return_value = {
+                "common": {
+                    "log_filename": "test.log"
+                },
+                "patch": {
+                    "start_issue": "2021-01-01",
+                    "end_issue": "2021-01-02",
+                    "patch_dir": "./patch_dir"
+                }
+            }
+
+            patch()
+
+            assert os.path.isdir('./patch_dir')
+            assert os.path.isdir('./patch_dir/issue_20210101/hospital-admissions')
+            assert os.path.isdir('./patch_dir/issue_20210102/hospital-admissions')
+
+            # Clean up the created directories after the test
+            shutil.rmtree(mock_read_params.return_value["patch"]["patch_dir"])