From 3f11f20b5291b3c1982e46695808240c1740e2c1 Mon Sep 17 00:00:00 2001 From: Vitor Avila Date: Wed, 29 Jan 2025 16:01:29 -0300 Subject: [PATCH 1/3] feat(import assets): Pass DB passwords in the import command --- .../cli/superset/sync/native/command.py | 30 ++- .../cli/superset/sync/native/command_test.py | 210 ++++++++++++++++-- 2 files changed, 221 insertions(+), 19 deletions(-) diff --git a/src/preset_cli/cli/superset/sync/native/command.py b/src/preset_cli/cli/superset/sync/native/command.py index 71d7012..fb20afc 100644 --- a/src/preset_cli/cli/superset/sync/native/command.py +++ b/src/preset_cli/cli/superset/sync/native/command.py @@ -226,6 +226,11 @@ def render_yaml(path: Path, env: Dict[str, Any]) -> Dict[str, Any]: "This way other asset types included get created but not overwritten." ), ) +@click.option( + "--db-password", + multiple=True, + help="Password for DB connections being imported (eg, uuid1=my_db_password)", +) @click.pass_context def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-branches ctx: click.core.Context, @@ -239,6 +244,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra load_env: bool = False, split: bool = False, continue_on_error: bool = False, + db_password: Tuple[str, ...] | None = None, ) -> None: """ Sync exported DBs/datasets/charts/dashboards to Superset. @@ -272,6 +278,12 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra if load_env: env["env"] = os.environ # type: ignore + pwds = ( + {kv.partition("=")[0]: kv.partition("=")[2] for kv in db_password} + if db_password + else {} + ) + # read all the YAML files configs: Dict[Path, AssetConfig] = {} queue = [root] @@ -306,7 +318,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra relative_path.parts[0] == "databases" and config["uuid"] not in existing_databases ): - prompt_for_passwords(relative_path, config) + add_password_to_config(relative_path, config, pwds) verify_db_connectivity(config) if relative_path.parts[0] == "datasets" and isinstance( config.get("params"), @@ -454,15 +466,23 @@ def verify_db_connectivity(config: Dict[str, Any]) -> None: _logger.debug(ex) -def prompt_for_passwords(path: Path, config: Dict[str, Any]) -> None: +def add_password_to_config( + path: Path, + config: Dict[str, Any], + pwds: Dict[str, Any], +) -> None: """ - Prompt user for masked passwords. + Add password passed in the command to the config. - Modify the config in place. + Prompt user for masked passwords for new connections if not provided. Modify + the config in place. """ uri = config["sqlalchemy_uri"] password = make_url(uri).password - if password == PASSWORD_MASK and config.get("password") is None: + + if config["uuid"] in pwds: + config["password"] = pwds[config["uuid"]] + elif password == PASSWORD_MASK and config.get("password") is None: config["password"] = getpass.getpass( f"Please provide the password for {path}: ", ) diff --git a/tests/cli/superset/sync/native/command_test.py b/tests/cli/superset/sync/native/command_test.py index 81d8c19..1effe15 100644 --- a/tests/cli/superset/sync/native/command_test.py +++ b/tests/cli/superset/sync/native/command_test.py @@ -23,33 +23,46 @@ from preset_cli.cli.superset.main import superset_cli from preset_cli.cli.superset.sync.native.command import ( ResourceType, + add_password_to_config, import_resources, import_resources_individually, load_user_modules, - prompt_for_passwords, raise_helper, verify_db_connectivity, ) from preset_cli.exceptions import ErrorLevel, ErrorPayload, SupersetError -def test_prompt_for_passwords(mocker: MockerFixture) -> None: +def test_add_password_to_config(mocker: MockerFixture) -> None: """ - Test ``prompt_for_passwords``. + Test ``add_password_to_config``. """ getpass = mocker.patch("preset_cli.cli.superset.sync.native.command.getpass") - config = {"sqlalchemy_uri": "postgresql://user:XXXXXXXXXX@host:5432/db"} + config = { + "sqlalchemy_uri": "postgresql://user:XXXXXXXXXX@host:5432/db", + "uuid": "uuid1", + } path = Path("/path/to/root/databases/psql.yaml") - prompt_for_passwords(path, config) + add_password_to_config(path, config, {}) getpass.getpass.assert_called_with(f"Please provide the password for {path}: ") config["password"] = "password123" getpass.reset_mock() - prompt_for_passwords(path, config) + add_password_to_config(path, config, {}) getpass.getpass.assert_not_called() + getpass.reset_mock() + add_password_to_config(path, config, {"uuid1": "password321"}) + getpass.getpass.assert_not_called() + # db_password takes precedence over config["password"] + assert config == { + "sqlalchemy_uri": "postgresql://user:XXXXXXXXXX@host:5432/db", + "uuid": "uuid1", + "password": "password321", + } + def test_import_resources(mocker: MockerFixture) -> None: """ @@ -418,8 +431,8 @@ def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> No client.get_databases.return_value = [] mocker.patch("preset_cli.cli.superset.sync.native.command.import_resources") mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth") - prompt_for_passwords = mocker.patch( - "preset_cli.cli.superset.sync.native.command.prompt_for_passwords", + add_password_to_config = mocker.patch( + "preset_cli.cli.superset.sync.native.command.add_password_to_config", ) runner = CliRunner() @@ -430,9 +443,9 @@ def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> No catch_exceptions=False, ) assert result.exit_code == 0 - prompt_for_passwords.assert_called() + add_password_to_config.assert_called() - prompt_for_passwords.reset_mock() + add_password_to_config.reset_mock() client.get_databases.return_value = [ {"uuid": "uuid1"}, ] @@ -442,7 +455,7 @@ def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> No catch_exceptions=False, ) assert result.exit_code == 0 - prompt_for_passwords.assert_not_called() + add_password_to_config.assert_not_called() client.get_uuids.assert_not_called() @@ -610,8 +623,8 @@ def test_native_legacy_instance(mocker: MockerFixture, fs: FakeFilesystem) -> No client.get_uuids.return_value = {1: "uuid1"} mocker.patch("preset_cli.cli.superset.sync.native.command.import_resources") mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth") - prompt_for_passwords = mocker.patch( - "preset_cli.cli.superset.sync.native.command.prompt_for_passwords", + add_password_to_config = mocker.patch( + "preset_cli.cli.superset.sync.native.command.add_password_to_config", ) runner = CliRunner() @@ -622,7 +635,7 @@ def test_native_legacy_instance(mocker: MockerFixture, fs: FakeFilesystem) -> No catch_exceptions=False, ) assert result.exit_code == 0 - prompt_for_passwords.assert_not_called() + add_password_to_config.assert_not_called() def test_load_user_modules(mocker: MockerFixture, fs: FakeFilesystem) -> None: @@ -2001,3 +2014,172 @@ def test_native_invalid_asset_type(mocker: MockerFixture, fs: FakeFilesystem) -> assert result.exit_code == 2 assert "Invalid value for '--asset-type'" in result.output + + +def test_native_with_db_passwords(mocker: MockerFixture, fs: FakeFilesystem) -> None: + # pylint: disable=line-too-long + """ + Test the ``native`` command while passing db passwords in the command. + """ + root = Path("/path/to/root") + fs.create_dir(root) + sqlalchemy_uri = { + "sqlalchemy_uri": "postgresql://user:XXXXXXXXXX@host:5432/db", + } + unmasked_uri = { + "sqlalchemy_uri": "postgresql://user:unmasked@host:5432/db", + } + + db_configs = { + "db_config_masked_no_password": { + "uuid": "uuid1", + **sqlalchemy_uri, + }, + "other_db_config_masked_no_password": { + "uuid": "uuid2", + **sqlalchemy_uri, + }, + "db_config_masked_with_password": { + "uuid": "uuid3", + "password": "directpwd!", + **sqlalchemy_uri, + }, + "other_db_config_masked_with_password": { + "uuid": "uuid4", + "password": "directpwd!again", + **sqlalchemy_uri, + }, + "db_config_unmasked": { + "uuid": "uuid5", + **unmasked_uri, + }, + "other_db_config_unmasked": { + "uuid": "uuid6", + **unmasked_uri, + }, + "db_config_unmasked_with_password": { + "uuid": "uuid7", + "password": "unmaskedpwd!", + **unmasked_uri, + }, + "final_db_config_unmasked_with_password": { + "uuid": "uuid8", + "password": "unmaskedpwd!again", + **unmasked_uri, + }, + } + + for file_name, content in db_configs.items(): + fs.create_file( + root / f"databases/{file_name}.yaml", + contents=yaml.dump(content), + ) + + SupersetClient = mocker.patch( + "preset_cli.cli.superset.sync.native.command.SupersetClient", + ) + client = SupersetClient() + client.get_databases.return_value = [] + import_resources = mocker.patch( + "preset_cli.cli.superset.sync.native.command.import_resources", + ) + mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth") + mocker.patch("preset_cli.cli.superset.sync.native.command.verify_db_connectivity") + getpass = mocker.patch("preset_cli.cli.superset.sync.native.command.getpass") + getpass.getpass.return_value = "pwd_from_prompt" + + runner = CliRunner() + result = runner.invoke( + superset_cli, + [ + "https://superset.example.org/", + "sync", + "native", + str(root), + "--db-password", + "uuid1=pwd_from_command=1", + "--db-password", + "uuid3=pwd_from_command=2", + "--db-password", + "uuid5=pwd_from_command=3", + "--db-password", + "uuid7=pwd_from_command=4", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + contents = { + "bundle/databases/db_config_masked_no_password.yaml": yaml.dump( + { + "uuid": "uuid1", + **sqlalchemy_uri, + "password": "pwd_from_command=1", + "is_managed_externally": False, + }, + ), + "bundle/databases/other_db_config_masked_no_password.yaml": yaml.dump( + { + "uuid": "uuid2", + **sqlalchemy_uri, + "password": "pwd_from_prompt", + "is_managed_externally": False, + }, + ), + "bundle/databases/db_config_masked_with_password.yaml": yaml.dump( + { + "uuid": "uuid3", + **sqlalchemy_uri, + "password": "pwd_from_command=2", + "is_managed_externally": False, + }, + ), + "bundle/databases/other_db_config_masked_with_password.yaml": yaml.dump( + { + "uuid": "uuid4", + **sqlalchemy_uri, + "password": "directpwd!again", + "is_managed_externally": False, + }, + ), + "bundle/databases/db_config_unmasked.yaml": yaml.dump( + { + "uuid": "uuid5", + **unmasked_uri, + "password": "pwd_from_command=3", + "is_managed_externally": False, + }, + ), + "bundle/databases/other_db_config_unmasked.yaml": yaml.dump( + { + "uuid": "uuid6", + **unmasked_uri, + "is_managed_externally": False, + }, + ), + "bundle/databases/db_config_unmasked_with_password.yaml": yaml.dump( + { + "uuid": "uuid7", + **unmasked_uri, + "password": "pwd_from_command=4", + "is_managed_externally": False, + }, + ), + "bundle/databases/final_db_config_unmasked_with_password.yaml": yaml.dump( + { + "uuid": "uuid8", + **unmasked_uri, + "password": "unmaskedpwd!again", + "is_managed_externally": False, + }, + ), + } + + import_resources.assert_called_once_with( + contents, + client, + False, + ResourceType.ASSET, + ) + getpass.getpass.assert_called_once_with( + "Please provide the password for databases/other_db_config_masked_no_password.yaml: ", + ) From 6cee934c32bb05f1216dc128e8a22f0f4b7bce0c Mon Sep 17 00:00:00 2001 From: Vitor Avila Date: Wed, 29 Jan 2025 16:07:06 -0300 Subject: [PATCH 2/3] Removing unused skip --- tests/cli/superset/sync/native/command_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/cli/superset/sync/native/command_test.py b/tests/cli/superset/sync/native/command_test.py index 1effe15..d9ff3e2 100644 --- a/tests/cli/superset/sync/native/command_test.py +++ b/tests/cli/superset/sync/native/command_test.py @@ -2017,7 +2017,6 @@ def test_native_invalid_asset_type(mocker: MockerFixture, fs: FakeFilesystem) -> def test_native_with_db_passwords(mocker: MockerFixture, fs: FakeFilesystem) -> None: - # pylint: disable=line-too-long """ Test the ``native`` command while passing db passwords in the command. """ From 7d969014b3548dd73fca633277dee935a6451471 Mon Sep 17 00:00:00 2001 From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com> Date: Wed, 29 Jan 2025 19:31:58 -0300 Subject: [PATCH 3/3] Implementing PR feedback Co-authored-by: Beto Dealmeida --- src/preset_cli/cli/superset/sync/native/command.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/preset_cli/cli/superset/sync/native/command.py b/src/preset_cli/cli/superset/sync/native/command.py index fb20afc..4356dfd 100644 --- a/src/preset_cli/cli/superset/sync/native/command.py +++ b/src/preset_cli/cli/superset/sync/native/command.py @@ -278,11 +278,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra if load_env: env["env"] = os.environ # type: ignore - pwds = ( - {kv.partition("=")[0]: kv.partition("=")[2] for kv in db_password} - if db_password - else {} - ) + pwds = dict(kv.split("=", 1) for kv in db_password or []) # read all the YAML files configs: Dict[Path, AssetConfig] = {}