feat: allow passing a JWT to superset

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)

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")