feat: retry imports on connection error

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

@@ -57,6 +57,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

@@ -13,7 +13,9 @@
 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
@@ -213,14 +215,17 @@ def import_resources_individually(
     related_configs: Dict[str, Dict[Path, AssetConfig]] = {}
     for resource_name, get_related_uuids in imports:
         for path, config in configs.items():
-            if path.parts[1] == resource_name:
-                asset_configs = {path: config}
-                for uuid in get_related_uuids(config):
-                    asset_configs.update(related_configs[uuid])
-                _logger.info("Importing %s", path.relative_to("bundle"))
-                contents = {str(k): yaml.dump(v) for k, v in asset_configs.items()}
-                import_resources(contents, client, overwrite)
-                related_configs[config["uuid"]] = asset_configs
+            if path.parts[1] != resource_name:
+                continue
+
+            asset_configs = {path: config}
+            for uuid in get_related_uuids(config):
+                asset_configs.update(related_configs[uuid])
+
+            _logger.info("Importing %s", path.relative_to("bundle"))
+            contents = {str(k): yaml.dump(v) for k, v in asset_configs.items()}
+            import_resources(contents, client, overwrite)
+            related_configs[config["uuid"]] = asset_configs
 
 
 def get_charts_uuids(config: AssetConfig) -> Iterator[str]:
@@ -267,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
@@ -20,6 +21,7 @@
 from preset_cli.cli.superset.main import superset_cli
 from preset_cli.cli.superset.sync.native.command import (
     import_resources,
+    import_resources_individually,
     load_user_modules,
     prompt_for_passwords,
     raise_helper,
@@ -586,3 +588,37 @@ def test_native_split(mocker: MockerFixture, fs: FakeFilesystem) -> 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()
+
+    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)
+
+    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)
+    assert str(excinfo.value) == "Connection aborted."