Use backoff

dev-requirements.txt

@@ -10,6 +10,7 @@ aiosignal==1.2.0
 appdirs==1.4.4
 async-timeout==4.0.2
 attrs==22.1.0
+backoff==2.2.1
 beautifulsoup4==4.11.1
 build==0.8.0
 certifi==2021.10.8

requirements.txt

@@ -10,6 +10,7 @@ aiosignal==1.2.0
 appdirs==1.4.4
 async-timeout==4.0.2
 attrs==22.1.0
+backoff==2.2.1
 beautifulsoup4==4.11.1
 certifi==2021.10.8
 charset-normalizer==2.0.12

setup.cfg

@@ -53,6 +53,7 @@ install_requires =
     Cython>=0.29.26
     PyYAML>=6.0
     appdirs>=1.4.4
+    backoff>=2.2.1 
     beautifulsoup4>=4.10.0
     click>=8.0.3
     jinja2>=3.0.3

src/preset_cli/cli/superset/sync/native/command.py

@@ -6,14 +6,16 @@
 import importlib.util
 import logging
 import os
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timezone
 from io import BytesIO
 from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, Iterator, Tuple
 from zipfile import ZipFile
 
+import backoff
 import click
+import requests
 import yaml
 from jinja2 import Template
 from sqlalchemy.engine import create_engine
@@ -194,8 +196,6 @@ def import_resources_individually(
     configs: Dict[Path, AssetConfig],
     client: SupersetClient,
     overwrite: bool,
-    retries: int = 3,
-    retry_wait: timedelta = timedelta(seconds=5),
 ) -> None:
     """
     Import contents individually.
@@ -224,25 +224,7 @@ def import_resources_individually(
 
             _logger.info("Importing %s", path.relative_to("bundle"))
             contents = {str(k): yaml.dump(v) for k, v in asset_configs.items()}
-
-            # use retries, since when importing assets individually there's a higher
-            # chance that a request will fail
-            for retry in range(retries + 1):
-                if retry > 0:
-                    _logger.info("Retrying (%d/%d)", retry, retries)
-                try:
-                    import_resources(contents, client, overwrite)
-                    break
-                except ConnectionError:
-                    if retry < retries:
-                        _logger.warning(
-                            "Failed to connect, will retry in %d seconds",
-                            retry_wait.total_seconds(),
-                        )
-            else:
-                _logger.error("Failed to connect, no more retries left... giving up")
-                raise Exception("Unable to connect")
-
+            import_resources(contents, client, overwrite)
             related_configs[config["uuid"]] = asset_configs
 
 
@@ -290,6 +272,13 @@ def prompt_for_passwords(path: Path, config: Dict[str, Any]) -> None:
         )
 
 
+@backoff.on_exception(
+    backoff.expo,
+    (requests.exceptions.ConnectionError, requests.exceptions.Timeout),
+    max_time=60,
+    max_tries=5,
+    logger=__name__,
+)
 def import_resources(
     contents: Dict[str, str],
     client: SupersetClient,

tests/cli/superset/sync/native/command_test.py

@@ -9,6 +9,7 @@
 from zipfile import ZipFile
 
 import pytest
+import requests
 import yaml
 from click.testing import CliRunner
 from freezegun import freeze_time
@@ -589,43 +590,35 @@ def test_native_split(mocker: MockerFixture, fs: FakeFilesystem) -> None:
     )
 
 
-def test_import_resources_individually_retries(mocker: MockerFixture) -> None:
+def test_import_resources_individually_retries(
+    mocker: MockerFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     """
     Test retries in ``import_resources_individually``.
     """
+    # prevent test from actually waiting between tries
+    monkeypatch.setattr("time.sleep", lambda x: None)
+
     client = mocker.MagicMock()
-    _logger = mocker.patch("preset_cli.cli.superset.sync.native.command._logger")
 
-    mocker.patch(
-        "preset_cli.cli.superset.sync.native.command.import_resources",
-        side_effect=[
-            ConnectionError("Connection aborted."),
-            ConnectionError("Connection aborted."),
-            None,
-        ],
-    )
+    client.import_zip.side_effect = [
+        requests.exceptions.ConnectionError("Connection aborted."),
+        requests.exceptions.ConnectionError("Connection aborted."),
+        None,
+    ]
     contents = {
         Path("bundle/databases/gsheets.yaml"): {"name": "my database", "uuid": "uuid1"},
     }
     import_resources_individually(contents, client, overwrite=True)
-    _logger.warning.assert_has_calls(
-        [
-            mocker.call("Failed to connect, will retry in %d seconds", 5),
-            mocker.call("Failed to connect, will retry in %d seconds", 5),
-        ],
-    )
 
-    mocker.patch(
-        "preset_cli.cli.superset.sync.native.command.import_resources",
-        side_effect=[
-            ConnectionError("Connection aborted."),
-            ConnectionError("Connection aborted."),
-            None,
-        ],
-    )
+    client.import_zip.side_effect = [
+        requests.exceptions.ConnectionError("Connection aborted."),
+        requests.exceptions.ConnectionError("Connection aborted."),
+        requests.exceptions.ConnectionError("Connection aborted."),
+        requests.exceptions.ConnectionError("Connection aborted."),
+        requests.exceptions.ConnectionError("Connection aborted."),
+    ]
     with pytest.raises(Exception) as excinfo:
-        import_resources_individually(contents, client, overwrite=True, retries=1)
-    _logger.error.assert_called_with(
-        "Failed to connect, no more retries left... giving up",
-    )
-    assert str(excinfo.value) == "Unable to connect"
+        import_resources_individually(contents, client, overwrite=True)
+    assert str(excinfo.value) == "Connection aborted."