feat: export resources by ID

src/preset_cli/cli/main.py

@@ -21,25 +21,11 @@
 )
 from preset_cli.auth.main import Auth
 from preset_cli.cli.superset.main import superset
-from preset_cli.lib import setup_logging
+from preset_cli.lib import setup_logging, split_comma
 
 _logger = logging.getLogger(__name__)
 
 
-def split_comma(  # pylint: disable=unused-argument
-    ctx: click.core.Context,
-    param: str,
-    value: Optional[str],
-) -> List[str]:
-    """
-    Split CLI option into multiple values.
-    """
-    if value is None:
-        return []
-
-    return [option.strip() for option in value.split(",")]
-
-
 def get_status_icon(status: str) -> str:
     """
     Return an icon (emoji) for a given status.

src/preset_cli/cli/superset/export.py

@@ -4,15 +4,15 @@
 
 from collections import defaultdict
 from pathlib import Path
-from typing import Tuple
+from typing import List, Set, Tuple
 from zipfile import ZipFile
 
 import click
 import yaml
 from yarl import URL
 
 from preset_cli.api.clients.superset import SupersetClient
-from preset_cli.lib import remove_root
+from preset_cli.lib import remove_root, split_comma
 
 JINJA2_OPEN_MARKER = "__JINJA2_OPEN__"
 JINJA2_CLOSE_MARKER = "__JINJA2_CLOSE__"
@@ -32,11 +32,35 @@
     help="Asset type",
     multiple=True,
 )
+@click.option(
+    "--database-ids",
+    callback=split_comma,
+    help="Comma separated list of database IDs to export",
+)
+@click.option(
+    "--dataset-ids",
+    callback=split_comma,
+    help="Comma separated list of dataset IDs to export",
+)
+@click.option(
+    "--chart-ids",
+    callback=split_comma,
+    help="Comma separated list of chart IDs to export",
+)
+@click.option(
+    "--dashboard-ids",
+    callback=split_comma,
+    help="Comma separated list of dashboard IDs to export",
+)
 @click.pass_context
-def export_assets(  # pylint: disable=too-many-locals
+def export_assets(  # pylint: disable=too-many-locals, too-many-arguments
     ctx: click.core.Context,
     directory: str,
     asset_type: Tuple[str, ...],
+    database_ids: List[str],
+    dataset_ids: List[str],
+    chart_ids: List[str],
+    dashboard_ids: List[str],
     overwrite: bool = False,
 ) -> None:
     """
@@ -47,14 +71,21 @@ def export_assets(  # pylint: disable=too-many-locals
     client = SupersetClient(url, auth)
     root = Path(directory)
     asset_types = set(asset_type)
+    ids = {
+        "database": {int(id_) for id_ in database_ids},
+        "dataset": {int(id_) for id_ in dataset_ids},
+        "chart": {int(id_) for id_ in chart_ids},
+        "dashboard": {int(id_) for id_ in dashboard_ids},
+    }
 
     for resource_name in ["database", "dataset", "chart", "dashboard"]:
         if not asset_types or resource_name in asset_types:
-            export_resource(resource_name, root, client, overwrite)
+            export_resource(resource_name, ids[resource_name], root, client, overwrite)
 
 
 def export_resource(
     resource_name: str,
+    requested_ids: Set[int],
     root: Path,
     client: SupersetClient,
     overwrite: bool,
@@ -63,7 +94,11 @@ def export_resource(
     Export a given resource and unzip it in a directory.
     """
     resources = client.get_resources(resource_name)
-    ids = [resource["id"] for resource in resources]
+    ids = [
+        resource["id"]
+        for resource in resources
+        if resource["id"] in requested_ids or not requested_ids
+    ]
     buf = client.export_zip(resource_name, ids)
 
     with ZipFile(buf) as bundle:

src/preset_cli/cli/superset/main.py

@@ -37,7 +37,7 @@
 )
 @click.option("--loglevel", default="INFO")
 @click.pass_context
-def superset_cli(
+def superset_cli(  # pylint: disable=too-many-arguments
     ctx: click.core.Context,
     instance: str,
     jwt_token: Optional[str],

src/preset_cli/lib.py

@@ -5,8 +5,9 @@
 import json
 import logging
 from pathlib import Path
-from typing import Any, Dict, List, cast
+from typing import Any, Dict, List, Optional, cast
 
+import click
 from requests import Response
 from rich.logging import RichHandler
 
@@ -96,3 +97,17 @@ def validate_response(response: Response) -> None:
 
     _logger.error(message)
     raise SupersetError(errors=errors)
+
+
+def split_comma(  # pylint: disable=unused-argument
+    ctx: click.core.Context,
+    param: str,
+    value: Optional[str],
+) -> List[str]:
+    """
+    Split CLI option into multiple values.
+    """
+    if value is None:
+        return []
+
+    return [option.strip() for option in value.split(",")]

tests/api/clients/preset_test.py

@@ -23,7 +23,8 @@ def test_preset_client_get_teams(mocker: MockerFixture, requests_mock: Mocker) -
     teams = client.get_teams()
     assert teams == [1, 2, 3]
     _logger.debug.assert_called_with(
-        "GET %s", URL("https://ws.preset.io/api/v1/teams/")
+        "GET %s",
+        URL("https://ws.preset.io/api/v1/teams/"),
     )
 
 

tests/cli/main_test.py

@@ -13,12 +13,8 @@
 from pytest_mock import MockerFixture
 from yarl import URL
 
-from preset_cli.cli.main import (
-    get_status_icon,
-    parse_selection,
-    preset_cli,
-    split_comma,
-)
+from preset_cli.cli.main import get_status_icon, parse_selection, preset_cli
+from preset_cli.lib import split_comma
 
 
 def test_split_comma(mocker: MockerFixture) -> None:

tests/cli/superset/export_test.py

@@ -69,12 +69,24 @@ def test_export_resource(
     client = mocker.MagicMock()
     client.export_zip.return_value = dataset_export
 
-    export_resource(resource_name="database", root=root, client=client, overwrite=False)
+    export_resource(
+        resource_name="database",
+        requested_ids=set(),
+        root=root,
+        client=client,
+        overwrite=False,
+    )
     with open(root / "databases/gsheets.yaml", encoding="utf-8") as input_:
         assert input_.read() == "database_name: GSheets\nsqlalchemy_uri: gsheets://\n"
 
     # check that Jinja2 was escaped
-    export_resource(resource_name="dataset", root=root, client=client, overwrite=False)
+    export_resource(
+        resource_name="dataset",
+        requested_ids=set(),
+        root=root,
+        client=client,
+        overwrite=False,
+    )
     with open(root / "datasets/gsheets/test.yaml", encoding="utf-8") as input_:
         assert yaml.load(input_.read(), Loader=yaml.SafeLoader) == {
             "table_name": "test",
@@ -104,10 +116,17 @@ def test_export_resource_overwrite(
     client = mocker.MagicMock()
     client.export_zip.return_value = dataset_export
 
-    export_resource(resource_name="database", root=root, client=client, overwrite=False)
+    export_resource(
+        resource_name="database",
+        requested_ids=set(),
+        root=root,
+        client=client,
+        overwrite=False,
+    )
     with pytest.raises(Exception) as excinfo:
         export_resource(
             resource_name="database",
+            requested_ids=set(),
             root=root,
             client=client,
             overwrite=False,
@@ -117,7 +136,13 @@ def test_export_resource_overwrite(
         "/path/to/root/databases/gsheets.yaml"
     )
 
-    export_resource(resource_name="database", root=root, client=client, overwrite=True)
+    export_resource(
+        resource_name="database",
+        requested_ids=set(),
+        root=root,
+        client=client,
+        overwrite=True,
+    )
 
 
 def test_export_assets(mocker: MockerFixture, fs: FakeFilesystem) -> None:
@@ -142,10 +167,46 @@ def test_export_assets(mocker: MockerFixture, fs: FakeFilesystem) -> None:
     assert result.exit_code == 0
     export_resource.assert_has_calls(
         [
-            mock.call("database", Path("/path/to/root"), client, False),
-            mock.call("dataset", Path("/path/to/root"), client, False),
-            mock.call("chart", Path("/path/to/root"), client, False),
-            mock.call("dashboard", Path("/path/to/root"), client, False),
+            mock.call("database", set(), Path("/path/to/root"), client, False),
+            mock.call("dataset", set(), Path("/path/to/root"), client, False),
+            mock.call("chart", set(), Path("/path/to/root"), client, False),
+            mock.call("dashboard", set(), Path("/path/to/root"), client, False),
+        ],
+    )
+
+
+def test_export_assets_by_id(mocker: MockerFixture, fs: FakeFilesystem) -> None:
+    """
+    Test the ``export_assets`` command.
+    """
+    # root must exist for command to succeed
+    root = Path("/path/to/root")
+    fs.create_dir(root)
+
+    SupersetClient = mocker.patch("preset_cli.cli.superset.export.SupersetClient")
+    client = SupersetClient()
+    export_resource = mocker.patch("preset_cli.cli.superset.export.export_resource")
+    mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth")
+
+    runner = CliRunner()
+    result = runner.invoke(
+        superset_cli,
+        [
+            "https://superset.example.org/",
+            "export",
+            "/path/to/root",
+            "--database-ids",
+            "1,2,3",
+        ],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    export_resource.assert_has_calls(
+        [
+            mock.call("database", {1, 2, 3}, Path("/path/to/root"), client, False),
+            mock.call("dataset", set(), Path("/path/to/root"), client, False),
+            mock.call("chart", set(), Path("/path/to/root"), client, False),
+            mock.call("dashboard", set(), Path("/path/to/root"), client, False),
         ],
     )
 
@@ -180,8 +241,8 @@ def test_export_assets_by_type(mocker: MockerFixture, fs: FakeFilesystem) -> Non
     assert result.exit_code == 0
     export_resource.assert_has_calls(
         [
-            mock.call("dataset", Path("/path/to/root"), client, False),
-            mock.call("dashboard", Path("/path/to/root"), client, False),
+            mock.call("dataset", set(), Path("/path/to/root"), client, False),
+            mock.call("dashboard", set(), Path("/path/to/root"), client, False),
         ],
     )
 
@@ -208,10 +269,10 @@ def test_export_with_custom_auth(mocker: MockerFixture, fs: FakeFilesystem) -> N
     assert result.exit_code == 0
     export_resource.assert_has_calls(
         [
-            mock.call("database", Path("/path/to/root"), client, False),
-            mock.call("dataset", Path("/path/to/root"), client, False),
-            mock.call("chart", Path("/path/to/root"), client, False),
-            mock.call("dashboard", Path("/path/to/root"), client, False),
+            mock.call("database", set(), Path("/path/to/root"), client, False),
+            mock.call("dataset", set(), Path("/path/to/root"), client, False),
+            mock.call("chart", set(), Path("/path/to/root"), client, False),
+            mock.call("dashboard", set(), Path("/path/to/root"), client, False),
         ],
     )
 

tests/cli/superset/main_test.py

@@ -135,9 +135,13 @@ def test_superset() -> None:
         == """Usage: superset export [OPTIONS] DIRECTORY
 
 Options:
-  --overwrite        Overwrite existing resources
-  --asset-type TEXT  Asset type
-  --help             Show this message and exit.
+  --overwrite           Overwrite existing resources
+  --asset-type TEXT     Asset type
+  --database-ids TEXT   Comma separated list of database IDs to export
+  --dataset-ids TEXT    Comma separated list of dataset IDs to export
+  --chart-ids TEXT      Comma separated list of chart IDs to export
+  --dashboard-ids TEXT  Comma separated list of dashboard IDs to export
+  --help                Show this message and exit.
 """
     )