From 7806a894aa43cb0f3e25b8bfa73af5858114537c Mon Sep 17 00:00:00 2001 From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com> Date: Wed, 16 Oct 2024 00:28:53 -0300 Subject: [PATCH 1/2] feat: Support importing specific asset types (#322) --- .../cli/superset/sync/native/command.py | 81 ++++++- .../cli/superset/sync/native/command_test.py | 203 +++++++++++++++++- 2 files changed, 267 insertions(+), 17 deletions(-) diff --git a/src/preset_cli/cli/superset/sync/native/command.py b/src/preset_cli/cli/superset/sync/native/command.py index 122c14e..adf401d 100644 --- a/src/preset_cli/cli/superset/sync/native/command.py +++ b/src/preset_cli/cli/superset/sync/native/command.py @@ -8,10 +8,11 @@ import logging import os from datetime import datetime, timezone +from enum import Enum from io import BytesIO from pathlib import Path from types import ModuleType -from typing import Any, Dict, Iterator, Set, Tuple +from typing import Any, Dict, Iterator, Optional, Set, Tuple from zipfile import ZipFile import backoff @@ -41,12 +42,57 @@ AssetConfig = Dict[str, Any] -resource_types = { - "chart": "Slice", - "dashboard": "Dashboard", - "database": "Database", - "dataset": "SqlaTable", -} +class ResourceType(Enum): + """ + ResourceType Enum. Used to identify asset type (and corresponding metadata). + """ + + def __new__( + cls, + resource_name: str, + metadata_type: Optional[str] = None, + ) -> "ResourceType": + """ + ResourceType Constructor. + """ + obj = object.__new__(cls) + obj._value_ = resource_name + obj._resource_name = resource_name # type:ignore + obj._metadata_type = metadata_type # type:ignore + return obj + + @property + def resource_name(self) -> str: + """ + Return the resource name for the asset type. + """ + return self._resource_name # type: ignore + + @property + def metadata_type(self) -> str: + """ + Return the metadata type for the asset type. + """ + return self._metadata_type # type: ignore + + CHART = ("chart", "Slice") + DASHBOARD = ("dashboard", "Dashboard") + DATABASE = ("database", "Database") + DATASET = ("dataset", "SqlaTable") + + +def normalize_to_enum( # pylint: disable=unused-argument + ctx: click.core.Context, + param: str, + value: Optional[str], +): + """ + Normalize the ``--asset-type`` option value and return the + corresponding ResourceType Enum. + """ + if value is None: + return None + return ResourceType(value.lower()) def load_user_modules(root: Path) -> Dict[str, ModuleType]: @@ -155,6 +201,15 @@ def render_yaml(path: Path, env: Dict[str, Any]) -> Dict[str, Any]: default=False, help="Split imports into individual assets", ) +@click.option( + "--asset-type", + type=click.Choice([rt.resource_name for rt in ResourceType], case_sensitive=False), + callback=normalize_to_enum, + help=( + "Specify an asset type to import resources using the type's endpoint. " + "This way other asset types included get created but not overwritten." + ), +) @click.pass_context def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-branches ctx: click.core.Context, @@ -166,6 +221,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra external_url_prefix: str = "", load_env: bool = False, split: bool = False, + asset_type: Optional[ResourceType] = None, ) -> None: """ Sync exported DBs/datasets/charts/dashboards to Superset. @@ -245,7 +301,7 @@ def native( # pylint: disable=too-many-locals, too-many-arguments, too-many-bra import_resources_individually(configs, client, overwrite) else: contents = {str(k): yaml.dump(v) for k, v in configs.items()} - import_resources(contents, client, overwrite) + import_resources(contents, client, overwrite, asset_type=asset_type) def import_resources_individually( @@ -257,7 +313,7 @@ def import_resources_individually( Import contents individually. This will first import all the databases, then import each dataset (together with the - database info, since it's needed), then charts, on so on. It helps troubleshoot + database info, since it's needed), then charts, and so on. It helps troubleshoot problematic exports and large imports. """ # store progress in case the import stops midway @@ -375,14 +431,17 @@ def import_resources( contents: Dict[str, str], client: SupersetClient, overwrite: bool, + asset_type: Optional[ResourceType] = None, ) -> None: """ Import a bundle of assets. """ + metadata_type = asset_type.metadata_type if asset_type else "assets" + resource_name = asset_type.resource_name if asset_type else "assets" contents["bundle/metadata.yaml"] = yaml.dump( dict( version="1.0.0", - type="assets", + type=metadata_type, timestamp=datetime.now(tz=timezone.utc).isoformat(), ), ) @@ -394,7 +453,7 @@ def import_resources( output.write(file_content.encode()) buf.seek(0) try: - client.import_zip("assets", buf, overwrite=overwrite) + client.import_zip(resource_name, buf, overwrite=overwrite) except SupersetError as ex: click.echo( click.style( diff --git a/tests/cli/superset/sync/native/command_test.py b/tests/cli/superset/sync/native/command_test.py index 9ababc6..fe05d93 100644 --- a/tests/cli/superset/sync/native/command_test.py +++ b/tests/cli/superset/sync/native/command_test.py @@ -21,6 +21,7 @@ from preset_cli.cli.superset.main import superset_cli from preset_cli.cli.superset.sync.native.command import ( + ResourceType, import_resources, import_resources_individually, load_user_modules, @@ -76,6 +77,42 @@ def test_import_resources(mocker: MockerFixture) -> None: assert bundle.read("bundle/databases/gsheets.yaml").decode() == "GSheets" +@pytest.mark.parametrize("resource_type", list(ResourceType)) +def test_import_resources_asset_types( + mocker: MockerFixture, + resource_type: ResourceType, +) -> None: + """ + Test ``import_resources`` when a resource_type value is specified. + """ + client = mocker.MagicMock() + + contents = {"bundle/databases/gsheets.yaml": "GSheets"} + with freeze_time("2022-01-01T00:00:00Z"): + import_resources( + contents=contents, + client=client, + overwrite=False, + asset_type=resource_type, + ) + + call = client.import_zip.mock_calls[0] + assert call.kwargs == {"overwrite": False} + + resource, buf = call.args + assert resource == resource_type.resource_name + with ZipFile(buf) as bundle: + assert bundle.namelist() == [ + "bundle/databases/gsheets.yaml", + "bundle/metadata.yaml", + ] + assert bundle.read("bundle/metadata.yaml").decode() == ( + "timestamp: '2022-01-01T00:00:00+00:00'\n" + f"type: {resource_type.metadata_type}\nversion: 1.0.0\n" + ) + assert bundle.read("bundle/databases/gsheets.yaml").decode() == "GSheets" + + def test_import_resources_overwrite_needed(mocker: MockerFixture) -> None: """ Test ``import_resources`` when an overwrite error is raised. @@ -290,7 +327,9 @@ def test_native(mocker: MockerFixture, fs: FakeFilesystem) -> None: ), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) client.get_uuids.assert_not_called() @@ -347,7 +386,9 @@ def test_native_params_as_str(mocker: MockerFixture, fs: FakeFilesystem) -> None }, ), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) client.get_uuids.assert_not_called() @@ -459,7 +500,9 @@ def test_native_load_env( }, ), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) client.get_uuids.assert_not_called() @@ -523,7 +566,9 @@ def test_native_external_url(mocker: MockerFixture, fs: FakeFilesystem) -> None: "bundle/databases/gsheets.yaml": yaml.dump(database_config), "bundle/datasets/gsheets/test.yaml": yaml.dump(dataset_config), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) client.get_uuids.assert_not_called() @@ -662,7 +707,9 @@ def test_template_in_environment(mocker: MockerFixture, fs: FakeFilesystem) -> N }, ), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) client.get_uuids.assert_not_called() @@ -1154,5 +1201,149 @@ def test_sync_native_jinja_templating_disabled( "bundle/databases/gsheets.yaml": yaml.dump(database_config), "bundle/datasets/gsheets/test.yaml": yaml.dump(dataset_config), } - import_resources.assert_has_calls([mock.call(contents, client, False)]) + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=None)], + ) + client.get_uuids.assert_not_called() + + +@pytest.mark.parametrize("resource_type", list(ResourceType)) +def test_native_asset_types( + mocker: MockerFixture, + fs: FakeFilesystem, + resource_type: ResourceType, +) -> None: + """ + Test the ``native`` command while specifying an asset type. + """ + root = Path("/path/to/root") + fs.create_dir(root) + database_config = { + "database_name": "GSheets", + "sqlalchemy_uri": "gsheets://", + "is_managed_externally": False, + "uuid": "uuid1", + } + dataset_config = {"table_name": "test", "is_managed_externally": False} + chart_config = { + "slice_name": "test", + "viz_type": "big_number_total", + "params": { + "datasource": "1__table", + "viz_type": "big_number_total", + "slice_id": 1, + "metric": { + "expressionType": "SQL", + "sqlExpression": "COUNT(*)", + "column": None, + "aggregate": None, + "datasourceWarning": False, + "hasCustomLabel": True, + "label": "custom_calculation", + "optionName": "metric_6aq7h4t8b3t_jbp2rak398o", + }, + "adhoc_filters": [], + "header_font_size": 0.4, + "subheader_font_size": 0.15, + "y_axis_format": "SMART_NUMBER", + "time_format": "smart_date", + "extra_form_data": {}, + "dashboards": [], + }, + "query_context": None, + "is_managed_externally": False, + } + dashboard_config = { + "dashboard_title": "Some dashboard", + "is_managed_externally": False, + "position": { + "DASHBOARD_VERSION_KEY": "v2", + "CHART-BVI44PWH": { + "type": "CHART", + "meta": { + "uuid": "3", + }, + }, + }, + "metadata": {}, + "uuid": "4", + } + + fs.create_file( + root / "databases/gsheets.yaml", + contents=yaml.dump(database_config), + ) + fs.create_file( + root / "datasets/gsheets/test.yaml", + contents=yaml.dump(dataset_config), + ) + fs.create_file( + root / "charts/test_01.yaml", + contents=yaml.dump(chart_config), + ) + fs.create_file( + root / "dashboards/dashboard.yaml", + contents=yaml.dump(dashboard_config), + ) + + 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") + + runner = CliRunner() + result = runner.invoke( + superset_cli, + [ + "https://superset.example.org/", + "sync", + "native", + str(root), + "--asset-type", + resource_type.resource_name, + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + contents = { + "bundle/charts/test_01.yaml": yaml.dump(chart_config), + "bundle/databases/gsheets.yaml": yaml.dump(database_config), + "bundle/datasets/gsheets/test.yaml": yaml.dump(dataset_config), + "bundle/dashboards/dashboard.yaml": yaml.dump(dashboard_config), + } + + import_resources.assert_has_calls( + [mock.call(contents, client, False, asset_type=resource_type)], + ) client.get_uuids.assert_not_called() + + +def test_native_invalid_asset_type(mocker: MockerFixture, fs: FakeFilesystem) -> None: + """ + Test the ``native`` command while specifying an invalid asset type. + """ + root = Path("/path/to/root") + fs.create_dir(root) + mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth") + + runner = CliRunner() + result = runner.invoke( + superset_cli, + [ + "https://superset.example.org/", + "sync", + "native", + str(root), + "--asset-type", + "datasource", + ], + catch_exceptions=False, + ) + + assert result.exit_code == 2 + assert "Invalid value for '--asset-type'" in result.output From 4db8130e1be47f5c2008643ff25812234b220d92 Mon Sep 17 00:00:00 2001 From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com> Date: Wed, 16 Oct 2024 16:36:33 -0300 Subject: [PATCH 2/2] chore: Release 0.2.21 (#323) --- CHANGELOG.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b3545b1..58c4a32 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,6 +5,13 @@ Changelog Next ==== +Version 0.2.21 - 2024-10-16 +=========================== + +- It's now possible to set export commands to generate YAML files with Unix EOL (`#316 `_). +- Calls to the Preset Manager API are now paginated (`#320 `_). +- It's now possible to trigger an import that creates missing dependencies but don't overwrite them (`#321 `_). + Version 0.2.20 - 2024-08-20 ===========================