diff --git a/src/preset_cli/cli/superset/main.py b/src/preset_cli/cli/superset/main.py index 039f4204..095adbb7 100644 --- a/src/preset_cli/cli/superset/main.py +++ b/src/preset_cli/cli/superset/main.py @@ -1,11 +1,12 @@ """ Main entry point for Superset commands. """ -from typing import Any +from typing import Any, Optional import click from yarl import URL +from preset_cli.auth.jwt import JWTAuth from preset_cli.auth.main import UsernamePasswordAuth from preset_cli.cli.superset.export import ( export_assets, @@ -22,6 +23,7 @@ @click.group() @click.argument("instance") +@click.option("--jwt-token", default=None, help="JWT token") @click.option("-u", "--username", default="admin", help="Username") @click.option( "-p", @@ -36,6 +38,7 @@ def superset_cli( ctx: click.core.Context, instance: str, + jwt_token: Optional[str], username: str = "admin", password: str = "admin", ): @@ -48,7 +51,10 @@ def superset_cli( # allow a custom authenticator to be passed via the context if "AUTH" not in ctx.obj: - ctx.obj["AUTH"] = UsernamePasswordAuth(URL(instance), username, password) + if jwt_token: + ctx.obj["AUTH"] = JWTAuth(jwt_token) + else: + ctx.obj["AUTH"] = UsernamePasswordAuth(URL(instance), username, password) superset_cli.add_command(sql) diff --git a/tests/cli/superset/main_test.py b/tests/cli/superset/main_test.py index 170428a9..deb75212 100644 --- a/tests/cli/superset/main_test.py +++ b/tests/cli/superset/main_test.py @@ -4,8 +4,9 @@ import click from click.testing import CliRunner +from pytest_mock import MockerFixture -from preset_cli.cli.superset.main import mutate_commands, superset +from preset_cli.cli.superset.main import mutate_commands, superset, superset_cli def test_mutate_commands() -> None: @@ -138,3 +139,20 @@ def test_superset() -> None: --help Show this message and exit. """ ) + + +def test_superset_jwt_auth(mocker: MockerFixture) -> None: + """ + Test passing a JWT to authenticate with Superset. + """ + # pylint: disable=invalid-name + JWTAuth = mocker.patch("preset_cli.cli.superset.main.JWTAuth") + + runner = CliRunner() + runner.invoke( + superset_cli, + ["--jwt-token=SECRET", "http://localhost:8088/", "export"], + catch_exceptions=False, + ) + + JWTAuth.assert_called_with("SECRET")