added logging and more clean up

Author: aysim319

claims_hosp/delphi_claims_hosp/backfill.py

@@ -22,7 +22,7 @@
 
 gmpr = GeoMapper()
 
-def store_backfill_file(claims_filepath, _end_date, backfill_dir):
+def store_backfill_file(claims_filepath, _end_date, backfill_dir, logger):
     """
     Store county level backfill data into backfill_dir.
 
@@ -57,6 +57,7 @@ def store_backfill_file(claims_filepath, _end_date, backfill_dir):
     backfilldata = backfilldata.loc[(backfilldata["time_value"] >= _start_date)
                                     & (~backfilldata["fips"].isnull()),
                                     selected_columns]
+    logger.info("Filtering backfill data", startdate=_start_date, enddate=_end_date)
 
     backfilldata["lag"] = [(_end_date - x).days for x in backfilldata["time_value"]]
     backfilldata["time_value"] = backfilldata.time_value.dt.strftime("%Y-%m-%d")
@@ -69,10 +70,15 @@ def store_backfill_file(claims_filepath, _end_date, backfill_dir):
         "state_id": "string"
     })
 
-    path = backfill_dir + \
-        "/claims_hosp_as_of_%s.parquet"%datetime.strftime(_end_date, "%Y%m%d")
+    filename = "claims_hosp_as_of_%s.parquet"%datetime.strftime(_end_date, "%Y%m%d")
+    path = f"{backfill_dir}/{filename}"
+
     # Store intermediate file into the backfill folder
-    backfilldata.to_parquet(path, index=False)
+    try:
+        backfilldata.to_parquet(path, index=False)
+        logger.info("Stored backfill data in parquet", filename=filename)
+    except:
+        logger.info("Failed to store backfill data in parquet", )
     return path
 
 
@@ -90,26 +96,29 @@ def merge_existing_backfill_files(backfill_dir, backfill_file, issue_date, logge
     backfill_dir : str
         specified path to store backfill files.
     backfill_file : str
+        specific file add to merged backfill file.
     """
     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{6}", filename)
-            if len(pattern) == 1:
-                file_month = datetime.strptime(pattern[0], "%Y%m")
-                start_date = file_month.replace(day=1)
-                end_date = (start_date + timedelta(days=32)).replace(day=1)
-                if issue_date >= start_date and issue_date < end_date:
+            # need to only match files with 6 digits for merged files
+            pattern = re.findall(r"_(\d{6,6})\.parquet", filename)
+            if pattern:
+                file_month = datetime.strptime(pattern[0], "%Y%m").replace(day=1)
+                end_date = (file_month + timedelta(days=32)).replace(day=1)
+                if issue_date >= file_month and issue_date < end_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"))
+        logger.info("Issue date has no matching merged files", issue_date=issue_date.strftime("%Y-%m-%d"))
         return
 
+    logger.info("Adding missing date to merged file", issue_date=issue_date, filename=backfill_file, merged_filename=file_name)
+
     # Start to merge files
     merge_file = f"{file_name.split('.')[0]}_after_merge.parquet"
     try:
@@ -139,8 +148,10 @@ def merge_backfill_file(backfill_dir, most_recent, logger, test_mode=False):
         specified path to store backfill files.
     test_mode: bool
     """
-    new_files = glob.glob(backfill_dir + "/claims_hosp_as_of_*")
+    previous_month = (most_recent.replace(day=1) - timedelta(days=1)).strftime("%Y%m")
+    new_files = glob.glob(backfill_dir + f"/claims_hosp_as_of_{previous_month}*")
     if len(new_files) == 0: # if no any daily file is stored
+        logger.info("No new files to merge; skipping merging")
         return
 
     def get_date(file_link):
@@ -155,7 +166,7 @@ def get_date(file_link):
         logger.info("Not a new month; skipping merging")
         return
 
-
+    logger.info(f"Merging files", start_date=date_list[0], end_date=date_list[-1])
     # Start to merge files
     pdList = []
     for fn in new_files:

claims_hosp/tests/test_backfill.py

@@ -9,6 +9,7 @@
 import pytest
 
 # first party
+from delphi_utils.logger import get_structured_logger
 from delphi_claims_hosp.config import Config, GeoConstants
 from delphi_claims_hosp.backfill import store_backfill_file, merge_backfill_file, merge_existing_backfill_files
 
@@ -25,133 +26,146 @@
 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:
 
-    def test_store_backfill_file(self):
-        dropdate = datetime(2020, 1, 1) 
+    def cleanup(self):
+        for file in glob.glob(f"{backfill_dir}/*.parquet"):
+            os.remove(file)
+
+    def test_store_backfill_file(self, caplog):
+        dropdate = datetime(2020, 1, 1)
         fn = "claims_hosp_as_of_20200101.parquet"
-        assert fn not in os.listdir(backfill_dir)
-       
+        caplog.set_level(logging.INFO)
+        logger = get_structured_logger()
+        num_rows = len(pd.read_csv(DATA_FILEPATH))
+
         # Store backfill file
-        store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir)
+        store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir, logger)
         assert fn in os.listdir(backfill_dir)
+        assert "Stored backfill data in parquet" in caplog.text
+
+
         fn = "claims_hosp_as_of_20200101.parquet"
         backfill_df = pd.read_parquet(backfill_dir + "/"+ fn, engine='pyarrow')
-        
+
         selected_columns = ['time_value', 'fips', 'state_id',
                         'num', 'den', 'lag', 'issue_date']
-        assert set(selected_columns) == set(backfill_df.columns)  
-        
-        os.remove(backfill_dir + "/" + fn)
-        assert fn not in os.listdir(backfill_dir)
+
+        assert set(selected_columns) == set(backfill_df.columns)
+        assert num_rows == len(backfill_df)
+
+        self.cleanup()
 
-    def test_merge_backfill_file(self):
+    def test_merge_backfill_file(self, caplog):
         fn = "claims_hosp_202006.parquet"
-        assert fn not in os.listdir(backfill_dir)
-        
+        caplog.set_level(logging.INFO)
+        logger = get_structured_logger()
+
         # Check when there is no daily file to merge.
         today = datetime(2020, 6, 14)
-        merge_backfill_file(backfill_dir, today, TEST_LOGGER,
+        merge_backfill_file(backfill_dir, today, logger,
                             test_mode=True)
         assert fn not in os.listdir(backfill_dir)
-        
-        # Generate backfill daily files     
+        assert "No new files to merge; skipping merging" in caplog.text
+
+
+        # Generate backfill daily files
         for d in range(11, 15):
-            dropdate = datetime(2020, 6, d)        
-            store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir)
-        
-        # Check when the merged file is not generated
+            dropdate = datetime(2020, 6, d)
+            store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir, logger)
+
         today = datetime(2020, 7, 1)
-        merge_backfill_file(backfill_dir, today, TEST_LOGGER,
+        merge_backfill_file(backfill_dir, today, logger,
                             test_mode=True)
+        assert "Merging files" in caplog.text
         assert fn in os.listdir(backfill_dir)
 
         # Read daily file
         new_files = glob.glob(backfill_dir + "/claims_hosp_as_of_*.parquet")
-        pdList = []        
+        pdList = []
         for file in new_files:
-            if "from" in file:
-                continue
             df = pd.read_parquet(file, engine='pyarrow')
             pdList.append(df)
             os.remove(file)
         new_files = glob.glob(backfill_dir + "/claims_hosp*.parquet")
         assert len(new_files) == 1
 
         expected = pd.concat(pdList).sort_values(["time_value", "fips"])
-        
+
         # Read the merged file
         merged = pd.read_parquet(backfill_dir + "/" + fn, engine='pyarrow')
-        
+
         assert set(expected.columns) == set(merged.columns)
         assert expected.shape[0] == merged.shape[0]
         assert expected.shape[1] == merged.shape[1]
-        
-        os.remove(backfill_dir + "/" + fn)
-        assert fn not in os.listdir(backfill_dir)
 
-    def test_merge_existing_backfill_files(self):
+        self.cleanup()
+
+    def test_merge_existing_backfill_files(self, caplog):
         issue_date = datetime(year=2020, month=6, day=13)
         issue_date_str = issue_date.strftime("%Y%m%d")
+        caplog.set_level(logging.INFO)
+        logger = get_structured_logger()
         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)
+                store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir, logger)
 
-            today = datetime(2020, 6, 14)
+            today = datetime(2020, 7, 1)
             # creating expected file
-            merge_backfill_file(backfill_dir, today, TEST_LOGGER,
+            merge_backfill_file(backfill_dir, today, logger,
                                 test_mode=True)
-            original = f"{backfill_dir}/claims_hosp_from_20200611_to_20200614.parquet"
+            original = f"{backfill_dir}/claims_hosp_202006.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,
-                                test_mode=True, check_nd=2)
+            merge_backfill_file(backfill_dir, today, logger,
+                                test_mode=True)
 
             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)
+        file_to_add = store_backfill_file(DATA_FILEPATH, issue_date, backfill_dir, logger)
+        merge_existing_backfill_files(backfill_dir, file_to_add, issue_date, logger)
+
+        assert "Adding missing date to merged file" in caplog.text
 
         expected = pd.read_parquet(f"{backfill_dir}/expected.parquet")
-        merged = pd.read_parquet(f"{backfill_dir}/claims_hosp_from_20200611_to_20200614.parquet")
+        merged = pd.read_parquet(f"{backfill_dir}/claims_hosp_202006.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)
+        check = pd.concat([merged, expected]).drop_duplicates(keep=False)
 
+        assert len(check) == 0
 
-    def test_merge_existing_backfill_files_no_call(self):
+        self.cleanup()
+
+
+    def test_merge_existing_backfill_files_no_call(self, caplog):
         issue_date = datetime(year=2020, month=5, day=20)
-        issue_date_str = issue_date.strftime("%Y%m%d")
+        caplog.set_level(logging.INFO)
+        logger = get_structured_logger()
         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)
+                store_backfill_file(DATA_FILEPATH, dropdate, backfill_dir, logger)
 
             today = datetime(2020, 6, 14)
             # creating expected file
-            merge_backfill_file(backfill_dir, today, TEST_LOGGER,
+            merge_backfill_file(backfill_dir, today, logger,
                                 test_mode=True)
 
         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)
+        file_to_add = store_backfill_file(DATA_FILEPATH, issue_date, backfill_dir, logger)
+        merge_existing_backfill_files(backfill_dir, file_to_add, issue_date, logger)
+        assert "Issue date has no matching merged files" in caplog.text
 
-        old_files = glob.glob(backfill_dir + "*.parquet")
-        for file in old_files:
-            os.remove(file)
+        self.cleanup()