Skip to content

Commit

Permalink
chore(sync native): Attempt to get the DB connection UUIDs via API (#284
Browse files Browse the repository at this point in the history 
)

* chore(sync native): Attempt to get the DB connection UUIDs via API
* Fixing CI issues
  • Loading branch information
@Vitor-Avila
Vitor-Avila authored Apr 16, 2024
1 parent 85e0590 commit b0fd1ce
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 10 deletions.
11 changes: 10 additions & 1 deletion src/preset_cli/cli/superset/sync/native/command.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,16 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra
base_url = URL(external_url_prefix) if external_url_prefix else None

# collecting existing database UUIDs so we know if we're creating or updating
existing_databases = {str(uuid) for uuid in client.get_uuids("database").values()}
# newer versions expose the DB UUID in the API response,
# olders only expose it via export
try:
existing_databases = {
db_connection["uuid"] for db_connection in client.get_databases()
}
except KeyError:
existing_databases = {
str(uuid) for uuid in client.get_uuids("database").values()
}

# env for Jinja2 templating
env = dict(pair.split("=", 1) for pair in option if "=" in pair) # type: ignore
Expand Down
72 changes: 63 additions & 9 deletions tests/cli/superset/sync/native/command_test.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""
Tests for the native import command.
"""
# pylint: disable=redefined-outer-name, invalid-name
# pylint: disable=redefined-outer-name, invalid-name, too-many-lines

import json
from pathlib import Path
Expand Down Expand Up @@ -227,7 +227,7 @@ def test_native(mocker: MockerFixture, fs: FakeFilesystem) -> None:
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand Down Expand Up @@ -291,6 +291,7 @@ def test_native(mocker: MockerFixture, fs: FakeFilesystem) -> None:
}

import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()


def test_native_params_as_str(mocker: MockerFixture, fs: FakeFilesystem) -> None:
Expand Down Expand Up @@ -323,7 +324,7 @@ def test_native_params_as_str(mocker: MockerFixture, fs: FakeFilesystem) -> None
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand All @@ -347,6 +348,7 @@ def test_native_params_as_str(mocker: MockerFixture, fs: FakeFilesystem) -> None
),
}
import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()


def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> None:
Expand All @@ -370,7 +372,7 @@ def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> No
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
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(
Expand All @@ -388,14 +390,17 @@ def test_native_password_prompt(mocker: MockerFixture, fs: FakeFilesystem) -> No
prompt_for_passwords.assert_called()

prompt_for_passwords.reset_mock()
client.get_uuids.return_value = {"1": "uuid1"}
client.get_databases.return_value = [
{"uuid": "uuid1"},
]
result = runner.invoke(
superset_cli,
["https://superset.example.org/", "sync", "native", str(root)],
catch_exceptions=False,
)
assert result.exit_code == 0
prompt_for_passwords.assert_not_called()
client.get_uuids.assert_not_called()


def test_native_load_env(
Expand Down Expand Up @@ -425,7 +430,7 @@ def test_native_load_env(
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand Down Expand Up @@ -455,6 +460,7 @@ def test_native_load_env(
),
}
import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()


def test_native_external_url(mocker: MockerFixture, fs: FakeFilesystem) -> None:
Expand Down Expand Up @@ -488,7 +494,7 @@ def test_native_external_url(mocker: MockerFixture, fs: FakeFilesystem) -> None:
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand Down Expand Up @@ -518,6 +524,52 @@ def test_native_external_url(mocker: MockerFixture, fs: FakeFilesystem) -> None:
"bundle/datasets/gsheets/test.yaml": yaml.dump(dataset_config),
}
import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()


def test_native_legacy_instance(mocker: MockerFixture, fs: FakeFilesystem) -> None:
"""
Test ``native`` command for legacy instances that don't expose the
DB connection ``uuid`` in the API response.
"""
root = Path("/path/to/root")
fs.create_dir(root)
database_config = {
"database_name": "Postgres",
"sqlalchemy_uri": "postgresql://user:XXXXXXXXXX@host:5432/dbname",
"is_managed_externally": False,
"uuid": "uuid1",
}
fs.create_file(
root / "databases/gsheets.yaml",
contents=yaml.dump(database_config),
)

SupersetClient = mocker.patch(
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_databases.return_value = [
{
"connection_name": "Test",
},
]
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",
)

runner = CliRunner()

result = runner.invoke(
superset_cli,
["https://superset.example.org/", "sync", "native", str(root)],
catch_exceptions=False,
)
assert result.exit_code == 0
prompt_for_passwords.assert_not_called()


def test_load_user_modules(mocker: MockerFixture, fs: FakeFilesystem) -> None:
Expand Down Expand Up @@ -586,7 +638,7 @@ def test_template_in_environment(mocker: MockerFixture, fs: FakeFilesystem) -> N
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand All @@ -611,6 +663,7 @@ def test_template_in_environment(mocker: MockerFixture, fs: FakeFilesystem) -> N
),
}
import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()


def test_verify_db_connectivity(mocker: MockerFixture) -> None:
Expand Down Expand Up @@ -941,7 +994,7 @@ def test_sync_native_jinja_templating_disabled(
"preset_cli.cli.superset.sync.native.command.SupersetClient",
)
client = SupersetClient()
client.get_uuids.return_value = {}
client.get_databases.return_value = []
import_resources = mocker.patch(
"preset_cli.cli.superset.sync.native.command.import_resources",
)
Expand All @@ -965,3 +1018,4 @@ def test_sync_native_jinja_templating_disabled(
"bundle/datasets/gsheets/test.yaml": yaml.dump(dataset_config),
}
import_resources.assert_has_calls([mock.call(contents, client, False)])
client.get_uuids.assert_not_called()

0 comments on commit b0fd1ce

Please sign in to comment.