From 31d3847bc019ea2c4df482569cfaa33d3265a0ab Mon Sep 17 00:00:00 2001
From: Kyle Burton <burtonk@uchicago.edu>
Date: Fri, 19 Apr 2024 17:54:38 -0500
Subject: [PATCH] feat: Add backoff to sftp connection

---
 fence/sync/sync_users.py           |  9 +++++++--
 tests/dbgap_sync/test_user_sync.py | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/fence/sync/sync_users.py b/fence/sync/sync_users.py
index 2803c4523..bcecff599 100644
--- a/fence/sync/sync_users.py
+++ b/fence/sync/sync_users.py
@@ -1,3 +1,4 @@
+import backoff
 import glob
 import jwt
 import os
@@ -42,7 +43,7 @@
 from fence.resources.google.access_utils import GoogleUpdateException
 from fence.sync import utils
 from fence.sync.passport_sync.ras_sync import RASVisa
-from fence.utils import get_SQLAlchemyDriver
+from fence.utils import get_SQLAlchemyDriver, DEFAULT_BACKOFF_SETTINGS
 
 
 def _format_policy_id(path, privilege):
@@ -429,13 +430,17 @@ def _get_from_sftp_with_proxy(self, server, path):
                     parameters.get("port", "unknown"),
                 )
             )
-            client.connect(**parameters)
+            self._connect_with_ssh(ssh_client=client, parameters=parameters)
             with client.open_sftp() as sftp:
                 download_dir(sftp, "./", path)
 
         if proxy:
             proxy.close()
 
+    @backoff.on_exception(backoff.expo, Exception, **DEFAULT_BACKOFF_SETTINGS)
+    def _connect_with_ssh(self, ssh_client, parameters):
+        ssh_client.connect(**parameters)
+
     def _get_from_ftp_with_proxy(self, server, path):
         """
         Download data from ftp sever to a local dir
diff --git a/tests/dbgap_sync/test_user_sync.py b/tests/dbgap_sync/test_user_sync.py
index 0525df99b..49c69e88a 100644
--- a/tests/dbgap_sync/test_user_sync.py
+++ b/tests/dbgap_sync/test_user_sync.py
@@ -11,6 +11,7 @@
 from fence.resources.google.access_utils import GoogleUpdateException
 from fence.config import config
 from fence.job.visa_update_cronjob import Visa_Token_Update
+from fence.utils import DEFAULT_BACKOFF_SETTINGS
 
 from tests.dbgap_sync.conftest import (
     get_test_encoded_decoded_visa_and_exp,
@@ -497,6 +498,21 @@ def test_sync_with_google_errors(syncer, monkeypatch):
     syncer._update_arborist.assert_called()
     syncer._update_authz_in_arborist.assert_called()
 
+@patch("fence.sync.sync_users.paramiko.SSHClient")
+@patch("os.makedirs")
+@patch("os.path.exists", return_value=False)
+@pytest.mark.parametrize("syncer", ["google", "cleversafe"], indirect=True)
+def test_sync_with_sftp_connection_errors(mock_path, mock_makedir, mock_ssh_client, syncer, monkeypatch):
+    """
+    Verifies that when there is an sftp connection error connection, that the connection is retried the max amount of
+    tries as configured by DEFAULT_BACKOFF_SETTINGS
+    """
+    monkeypatch.setattr(syncer, "is_sync_from_dbgap_server", True)
+    mock_ssh_client.return_value.__enter__.return_value.connect.side_effect = Exception("Authentication timed out")
+    # usersync System Exits if any exception is raised during download.
+    with pytest.raises(SystemExit):
+        syncer.sync()
+    assert mock_ssh_client.return_value.__enter__.return_value.connect.call_count == DEFAULT_BACKOFF_SETTINGS['max_tries']
 
 @pytest.mark.parametrize("syncer", ["google", "cleversafe"], indirect=True)
 def test_sync_from_files(syncer, db_session, storage_client):