feat: check DB connectivity before import

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

@@ -4,6 +4,7 @@
 
 import getpass
 import importlib.util
+import logging
 import os
 from datetime import datetime, timezone
 from io import BytesIO
@@ -15,12 +16,15 @@
 import click
 import yaml
 from jinja2 import Template
+from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.url import make_url
 from yarl import URL
 
 from preset_cli.api.clients.superset import SupersetClient
 from preset_cli.exceptions import SupersetError
 
+_logger = logging.getLogger(__name__)
+
 YAML_EXTENSIONS = {".yaml", ".yml"}
 ASSET_DIRECTORIES = {"databases", "datasets", "charts", "dashboards"}
 
@@ -130,16 +134,16 @@ def native(  # pylint: disable=too-many-locals, too-many-arguments
                 env["filepath"] = path_name
                 template = Template(input_.read())
                 content = template.render(**env)
-
-                # mark resource as being managed externally
                 config = yaml.load(content, Loader=yaml.SafeLoader)
+
                 config["is_managed_externally"] = disallow_edits
                 if base_url:
                     config["external_url"] = str(
                         base_url / str(relative_path),
                     )
                 if relative_path.parts[0] == "databases":
                     prompt_for_passwords(relative_path, config)
+                    verify_db_connectivity(config)
 
                 contents[str("bundle" / relative_path)] = yaml.safe_dump(config)
 
@@ -148,6 +152,23 @@ def native(  # pylint: disable=too-many-locals, too-many-arguments
         import_resource(resource, contents, client, overwrite)
 
 
+def verify_db_connectivity(config: Dict[str, Any]) -> None:
+    """
+    Test if we can connect to a given database.
+    """
+    uri = make_url(config["sqlalchemy_uri"])
+    if config.get("password"):
+        uri = uri.set(password=config["password"])
+
+    try:
+        engine = create_engine(uri)
+        raw_connection = engine.raw_connection()
+        engine.dialect.do_ping(raw_connection)
+    except Exception as ex:  # pylint: disable=broad-except
+        _logger.warning("Cannot connect to database %s", uri)
+        _logger.debug(ex)
+
+
 def prompt_for_passwords(path: Path, config: Dict[str, Any]) -> None:
     """
     Prompt user for masked passwords.

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

@@ -15,13 +15,15 @@
 from jinja2 import Template
 from pyfakefs.fake_filesystem import FakeFilesystem
 from pytest_mock import MockerFixture
+from sqlalchemy.engine.url import URL
 
 from preset_cli.cli.superset.main import superset_cli
 from preset_cli.cli.superset.sync.native.command import (
     import_resource,
     load_user_modules,
     prompt_for_passwords,
     raise_helper,
+    verify_db_connectivity,
 )
 from preset_cli.exceptions import ErrorLevel, ErrorPayload, SupersetError
 
@@ -426,3 +428,76 @@ def test_template_in_environment(mocker: MockerFixture, fs: FakeFilesystem) -> N
             mock.call("dashboard", contents, client, False),
         ],
     )
+
+
+def test_verify_db_connectivity(mocker: MockerFixture) -> None:
+    """
+    Test ``verify_db_connectivity``.
+    """
+    create_engine = mocker.patch(
+        "preset_cli.cli.superset.sync.native.command.create_engine",
+    )
+
+    config = {
+        "sqlalchemy_uri": "postgresql://username:XXXXXXXXXX@localhost:5432/examples",
+        "password": "SECRET",
+    }
+    verify_db_connectivity(config)
+
+    create_engine.assert_called_with(
+        URL(
+            "postgresql",
+            username="username",
+            password="SECRET",
+            host="localhost",
+            port=5432,
+            database="examples",
+        ),
+    )
+
+
+def test_verify_db_connectivity_no_password(mocker: MockerFixture) -> None:
+    """
+    Test ``verify_db_connectivity`` without passwords.
+    """
+    create_engine = mocker.patch(
+        "preset_cli.cli.superset.sync.native.command.create_engine",
+    )
+
+    config = {
+        "sqlalchemy_uri": "gsheets://",
+    }
+    verify_db_connectivity(config)
+
+    create_engine.assert_called_with(
+        URL("gsheets"),
+    )
+
+
+def test_verify_db_connectivity_error(mocker: MockerFixture) -> None:
+    """
+    Test ``verify_db_connectivity`` errors.
+    """
+    _logger = mocker.patch("preset_cli.cli.superset.sync.native.command._logger")
+    mocker.patch(
+        "preset_cli.cli.superset.sync.native.command.create_engine",
+        side_effect=Exception("Unable to connect"),
+    )
+
+    config = {
+        "sqlalchemy_uri": "postgresql://username:XXXXXXXXXX@localhost:5432/examples",
+        "password": "SECRET",
+    }
+    verify_db_connectivity(config)
+
+    _logger.warning.assert_called_with(
+        "Cannot connect to database %s",
+        URL(
+            "postgresql",
+            username="username",
+            password="SECRET",
+            host="localhost",
+            port=5432,
+            database="examples",
+        ),
+    )