diff --git a/src/preset_cli/cli/superset/sync/native/command.py b/src/preset_cli/cli/superset/sync/native/command.py index 1cb1258c..a994254b 100644 --- a/src/preset_cli/cli/superset/sync/native/command.py +++ b/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,9 +134,8 @@ 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( @@ -140,6 +143,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments ) 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,22 @@ 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) + engine.execute("SELECT 1").scalar() + 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. diff --git a/tests/cli/superset/sync/native/command_test.py b/tests/cli/superset/sync/native/command_test.py index 1b5f2ce5..831cf0e6 100644 --- a/tests/cli/superset/sync/native/command_test.py +++ b/tests/cli/superset/sync/native/command_test.py @@ -15,6 +15,7 @@ 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 ( @@ -22,6 +23,7 @@ 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", + ), + )