change source host to params + add if logger before use in create_backup_csv

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/delphi_nssp/patch.py

@@ -15,6 +15,7 @@
     ...
   },
   "patch": {
+    "source_host": "prod.server.edu",
     "source_dir": "delphi/covidcast-indicators/nssp/source_data",
     "user": "username",
     "patch_dir": "delphi/covidcast-indicators/nssp/AprilPatch",
@@ -26,6 +27,7 @@
 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
@@ -72,7 +74,7 @@ def good_patch_config(params, logger):
         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"]
+        required_patch_keys = ["start_issue", "end_issue", "patch_dir", "source_dir", "source_host"]
 
         source_dir = params["patch"]["source_dir"]
         if not path.isdir(source_dir) or not listdir(source_dir):

nssp/delphi_nssp/pull.py

@@ -42,7 +42,7 @@ def get_source_data(params, logger):
     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"
+    host = params["patch"]["source_host"]
     user = params["patch"]["user"]
     ssh.connect(host, username=user)
 
@@ -51,29 +51,28 @@ def get_source_data(params, logger):
     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)
+    with ssh.open_sftp() as sftp:
         try:
-            sftp.stat(remote_file_name)
+            sftp.stat(params["common"]["backup_dir"])
         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()
+            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
 
     if num_files_transferred == 0:
         logger.error("No source data was transferred. Check the source backup server for potential issues.")

nssp/tests/test_patch.py

@@ -28,6 +28,7 @@ def test_config_missing_end_issue_in_patch_section(self, mock_logger):
             },
             "patch": {
                 "source_dir": "./does_not_exist",
+                "source_host": "prod.server.edu",
                 "user": "user",
                 "patch_dir": "dir",
                 "start_issue": "2024-04-21",
@@ -47,6 +48,7 @@ def test_config_invalid_start_issue(self, mock_logger):
             },
             "patch": {
                 "source_dir": "./does_not_exist",
+                "source_host": "prod.server.edu",
                 "user": "user",
                 "patch_dir": "dir",
                 "start_issue": "01-01-2024",
@@ -71,6 +73,7 @@ def test_config_start_issue_after_end_issue(self, mock_logger):
                 "start_issue": "2024-04-22",
                 "end_issue": "2024-04-21",
                 "source_dir": "./does_not_exist",
+                "source_host": "prod.server.edu"
             }
         }
         assert not good_patch_config(patch_config, mock_logger)
@@ -88,7 +91,8 @@ def test_config_all_valid_configurations(self, mock_logger):
                 "patch_dir": "dir",
                 "start_issue": "2024-04-21",
                 "end_issue": "2024-04-22",
-                "source_dir": "./does_not_exist"
+                "source_dir": "./does_not_exist",
+                "source_host": "prod.server.edu"
             }
         }
         assert good_patch_config(patch_config, mock_logger)
@@ -106,7 +110,8 @@ def test_config_user_param(self, mock_logger):
                 "patch_dir": "dir",
                 "start_issue": "2024-04-21",
                 "end_issue": "2024-04-22",
-                "source_dir": "./source_dir"
+                "source_dir": "./source_dir",
+                "source_host": "prod.server.edu"
             }
         }
         assert good_patch_config(patch_config, mock_logger)
@@ -157,7 +162,8 @@ def test_patch_from_local_source(self, mock_read_params, mock_get_structured_log
                 "start_issue": "2021-01-01",
                 "end_issue": "2021-01-16",
                 "patch_dir": "./patch_dir",
-                "source_dir": "./source_dir"
+                "source_dir": "./source_dir",
+                "source_host": "prod.server.edu"
             }
         }
         patch()
@@ -180,6 +186,7 @@ def test_patch_download_remote_source(self, mock_read_params, mock_get_structure
                 "start_issue": "2021-01-01",
                 "end_issue": "2021-01-16",
                 "patch_dir": "./patch_dir",
+                "source_host": "prod.server.edu",
                 "source_dir": "./does_not_exist"
             }
         }
@@ -205,7 +212,8 @@ def test_patch_confirm_dir_structure_created(self, mock_read_params, mock_get_st
                 "start_issue": "2021-01-01",
                 "end_issue": "2021-01-16",
                 "patch_dir": "./patch_dir",
-                "source_dir": "./source_dir"
+                "source_dir": "./source_dir",
+                "source_host": "prod.server.edu",
             }
         }
         patch()
@@ -243,7 +251,8 @@ def test_full_patch_code(self, mock_read_params, mock_get_structured_logger):
                 "start_issue": "2021-01-01",
                 "end_issue": "2021-01-16",
                 "patch_dir": "./patch_dir",
-                "source_dir": "./source_dir"
+                "source_dir": "./source_dir",
+                "source_host": "prod.server.edu"
             }
         }
 

nssp/tests/test_pull.py

@@ -29,6 +29,7 @@ def test_get_source_data(self):
         params = {
             "patch": {
                 "source_dir": "test_source_dir",
+                "source_host": "prod.server.edu",
                 "user": "test_user",
                 "start_issue": "2023-01-01",
                 "end_issue": "2023-01-03",
@@ -48,7 +49,6 @@ def test_get_source_data(self):
         with patch("paramiko.SSHClient", return_value=mock_ssh):
             get_source_data(params, logger)
 
-        mock_ssh.connect.assert_called_once_with("delphi-master-prod-01.delphi.cmu.edu", username="test_user")
         mock_sftp.chdir.assert_called_once_with("/test_backup_dir")
         assert mock_sftp.get.call_count == 3