tinybird 0.0.1.dev0__py3-none-any.whl

tinybird/tb_cli_modules/token.py

@@ -0,0 +1,349 @@
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
+
+import click
+import pyperclip
+from click import Context
+from humanfriendly import parse_timespan
+
+from tinybird.client import AuthNoTokenException, TinyB
+from tinybird.feedback_manager import FeedbackManager
+from tinybird.tb_cli_modules.cli import cli
+from tinybird.tb_cli_modules.common import (
+    DoesNotExistException,
+    coro,
+    echo_safe_humanfriendly_tables_format_smart_table,
+)
+from tinybird.tb_cli_modules.exceptions import CLITokenException
+
+
+@cli.group()
+@click.pass_context
+def token(ctx: Context) -> None:
+    """Token commands."""
+
+
+@token.command(name="ls")
+@click.option("--match", default=None, help="Retrieve any token matching the pattern. eg --match _test")
+@click.pass_context
+@coro
+async def token_ls(
+    ctx: Context,
+    match: Optional[str] = None,
+) -> None:
+    """List Static Tokens."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+
+    try:
+        tokens = await client.token_list(match)
+        columns = ["id", "name", "description"]
+        table = list(map(lambda token: [token.get(key, "") for key in columns], tokens))
+
+        click.echo(FeedbackManager.info_tokens())
+        echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
+        click.echo("\n")
+    except AuthNoTokenException:
+        raise
+    except Exception as e:
+        raise CLITokenException(FeedbackManager.error_exception(error=e))
+
+
+@token.command(name="rm")
+@click.argument("token_id")
+@click.option("--yes", is_flag=True, default=False, help="Do not ask for confirmation")
+@click.pass_context
+@coro
+async def token_rm(ctx: Context, token_id: str, yes: bool) -> None:
+    """Remove a static token."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+    if yes or click.confirm(FeedbackManager.warning_confirm_delete_token(token=token_id)):
+        try:
+            await client.token_delete(token_id)
+        except AuthNoTokenException:
+            raise
+        except DoesNotExistException:
+            raise CLITokenException(FeedbackManager.error_token_does_not_exist(token_id=token_id))
+        except Exception as e:
+            raise CLITokenException(FeedbackManager.error_exception(error=e))
+        click.echo(FeedbackManager.success_delete_token(token=token_id))
+
+
+@token.command(name="refresh")
+@click.argument("token_id")
+@click.option("--yes", is_flag=True, default=False, help="Do not ask for confirmation")
+@click.pass_context
+@coro
+async def token_refresh(ctx: Context, token_id: str, yes: bool) -> None:
+    """Refresh a Static Token."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+    if yes or click.confirm(FeedbackManager.warning_confirm_refresh_token(token=token_id)):
+        try:
+            await client.token_refresh(token_id)
+        except AuthNoTokenException:
+            raise
+        except DoesNotExistException:
+            raise CLITokenException(FeedbackManager.error_token_does_not_exist(token_id=token_id))
+        except Exception as e:
+            raise CLITokenException(FeedbackManager.error_exception(error=e))
+        click.echo(FeedbackManager.success_refresh_token(token=token_id))
+
+
+@token.command(name="scopes")
+@click.argument("token_id")
+@click.pass_context
+@coro
+async def token_scopes(ctx: Context, token_id: str) -> None:
+    """List Static Token scopes."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+
+    try:
+        scopes = await client.token_scopes(token_id)
+        columns = ["type", "resource", "filter"]
+        table = list(map(lambda scope: [scope.get(key, "") for key in columns], scopes))
+        click.echo(FeedbackManager.info_token_scopes(token=token_id))
+        echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
+        click.echo("\n")
+    except AuthNoTokenException:
+        raise
+    except Exception as e:
+        raise CLITokenException(FeedbackManager.error_exception(error=e))
+
+
+@token.command(name="copy")
+@click.argument("token_id")
+@click.pass_context
+@coro
+async def token_copy(ctx: Context, token_id: str) -> None:
+    """Copy a Static Token."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+
+    try:
+        token = await client.token_get(token_id)
+        pyperclip.copy(token["token"].strip())
+    except AuthNoTokenException:
+        raise
+    except DoesNotExistException:
+        raise CLITokenException(FeedbackManager.error_token_does_not_exist(token_id=token_id))
+    except Exception as e:
+        raise CLITokenException(FeedbackManager.error_exception(error=e))
+    click.echo(FeedbackManager.success_copy_token(token=token_id))
+
+
+def parse_ttl(ctx, param, value):
+    if value is None:
+        return None
+    try:
+        seconds = parse_timespan(value)
+        return timedelta(seconds=seconds)
+    except ValueError:
+        raise click.BadParameter(f"Invalid time to live format: {value}")
+
+
+def parse_fixed_params(fixed_params_list):
+    parsed_params = []
+    for fixed_param in fixed_params_list:
+        param_dict = {}
+        for param in fixed_param.split(","):
+            key, value = param.split("=")
+            param_dict[key] = value
+        parsed_params.append(param_dict)
+    return parsed_params
+
+
+@token.group()
+@click.pass_context
+def create(ctx: Context) -> None:
+    """Token creation commands.
+
+    You can create two types of tokens: JWT or Static.
+
+    * JWT tokens have a TTL and can only have the PIPES:READ scope.Their main use case is allow your users to call your endpoints without exposing your API key.
+
+    * Static Tokens do not have a TTL and can have any valid scope (DATASOURCES:READ, DATASOURCES:APPEND, DATASOURCES:CREATE, DATASOURCES:DROP, PIPES:CREATE, PIPES:READ, PIPES:DROP).
+
+    Examples:
+
+    tb token create jwt my_jwt_token --ttl 1h --scope PIPES:READ --resource my_pipe
+
+    tb token create static my_static_token --scope PIPES:READ --resource my_pipe
+
+    tb token create static my_static_token --scope DATASOURCES:READ --resource my_datasource
+
+    tb token create static my_static_token --scope DATASOURCES:READ --resource my_datasource --filters "column_name=value"
+
+    """
+
+
+@create.command(name="jwt")
+@click.argument("name")
+@click.option("--ttl", type=str, callback=parse_ttl, required=True, help="Time to live (e.g., '1h', '30min', '1d')")
+@click.option(
+    "--scope",
+    multiple=True,
+    type=click.Choice(["PIPES:READ"]),
+    required=True,
+    help="Scope of the token (only PIPES:READ is allowed for JWT tokens)",
+)
+@click.option("--resource", multiple=True, required=True, help="Resource associated with the scope")
+@click.option(
+    "--fixed-params", multiple=True, help="Fixed parameters in key=value format, multiple values separated by commas"
+)
+@click.pass_context
+@coro
+async def create_jwt_token(ctx: Context, name: str, ttl: timedelta, scope, resource, fixed_params) -> None:
+    """Create a JWT token with a TTL specify."""
+
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+
+    expiration_time = int((ttl + datetime.now(timezone.utc)).timestamp())
+    if len(scope) != len(resource):
+        raise CLITokenException(FeedbackManager.error_number_of_scopes_and_resources_mismatch())
+
+    # Ensure the number of fixed-params does not exceed the number of scope/resource pairs
+    if fixed_params and len(fixed_params) > len(scope):
+        raise CLITokenException(FeedbackManager.error_number_of_fixed_params_and_resources_mismatch())
+
+    # Parse fixed params
+    parsed_fixed_params = parse_fixed_params(fixed_params) if fixed_params else []
+
+    # Create a list of fixed params for each scope/resource pair, defaulting to empty dict if not provided
+    fixed_params_list: List[Dict[str, Any]] = [{}] * len(scope)
+    for i, params in enumerate(parsed_fixed_params):
+        fixed_params_list[i] = params
+
+    scopes = []
+    for sc, res, fparams in zip(scope, resource, fixed_params_list):
+        scopes.append(
+            {
+                "type": sc,
+                "resource": res,
+                "fixed_params": fparams,
+            }
+        )
+
+    try:
+        response = await client.create_jwt_token(name, expiration_time, scopes)
+    except AuthNoTokenException:
+        raise
+    except Exception as e:
+        raise CLITokenException(FeedbackManager.error_exception(error=e))
+
+    click.echo("The token has been generated successfully.")
+    click.echo(
+        f"The token will expire at: {datetime.fromtimestamp(expiration_time).strftime('%Y-%m-%d %H:%M:%S')} UTC "
+    )
+    click.echo(f"The token is: {response['token']}")
+
+
+# Valid scopes for Static Tokens
+valid_scopes = [
+    "DATASOURCES:READ",
+    "DATASOURCES:APPEND",
+    "DATASOURCES:CREATE",
+    "DATASOURCES:DROP",
+    "PIPES:CREATE",
+    "PIPES:READ",
+    "PIPES:DROP",
+]
+
+
+# As we are passing dynamic options to the command, we need to create a custom class to handle the help message
+class DynamicOptionsCommand(click.Command):
+    def get_help(self, ctx):
+        # Usage
+        usage = "Usage: tb token create static [OPTIONS] NAME\n\n"
+        dynamic_options_help = usage
+
+        # Description
+        dynamic_options_help += "  Create a Static Token that will live forever.\n\n"
+
+        # Options
+        dynamic_options_help += "Options:\n"
+        dynamic_options_help += f"  --scope [{','.join(valid_scopes)}]   Scope for the token [Required]\n"
+        dynamic_options_help += "  --resource TEXT    Resource you want to associate the scope with\n"
+        dynamic_options_help += "  --filter TEXT  SQL condition used to filter the values when calling with this token (eg. --filter=value > 0) \n"
+        dynamic_options_help += "  -h, --help            Show this message and exit.\n"
+
+        return dynamic_options_help
+
+
+@create.command(
+    name="static", context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=DynamicOptionsCommand
+)
+@click.argument("name")
+@click.pass_context
+@coro
+async def create_static_token(ctx, name: str):
+    """Create a Static Token."""
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    client: TinyB = obj["client"]
+
+    args = ctx.args
+    scopes: List[Dict[str, str]] = []
+    current_scope = None
+
+    # We parse the arguments to get the scopes, resources and filters
+    # The arguments should be in the format --scope <scope> --resource <resource> --filter <filter>
+    i = 0
+    while i < len(args):
+        if args[i] == "--scope":
+            if current_scope:
+                scopes.append(current_scope)
+                current_scope = {}
+            current_scope = {"scope": args[i + 1]}
+            i += 2
+        elif args[i] == "--resource":
+            if current_scope is None:
+                raise click.BadParameter("Resource must follow a scope")
+            if "resource" in current_scope:
+                raise click.BadParameter(
+                    "Resource already defined for this scope. The format is --scope <scope> --resource <resource> --filter <filter>"
+                )
+            current_scope["resource"] = args[i + 1]
+            i += 2
+        elif args[i] == "--filter":
+            if current_scope is None:
+                raise click.BadParameter("Filter must follow a scope")
+            if "filter" in current_scope:
+                raise click.BadParameter(
+                    "Filter already defined for this scope. The format is --scope <scope> --resource <resource> --filter <filter>"
+                )
+            current_scope["filter"] = args[i + 1]
+            i += 2
+        else:
+            raise click.BadParameter(f"Unknown parameter {args[i]}")
+
+    if current_scope:
+        scopes.append(current_scope)
+
+    # Parse the scopes like `SCOPE:RESOURCE:FILTER` or `SCOPE:RESOURCE` or `SCOPE` as that's what the API expsects
+    scoped_parsed: List[str] = []
+    for scope in scopes:
+        if scope.get("resource") and scope.get("filter"):
+            scoped_parsed.append(f"{scope.get('scope')}:{scope.get('resource')}:{scope.get('filter')}")
+        elif scope.get("resource"):
+            scoped_parsed.append(f"{scope.get('scope')}:{scope.get('resource')}")
+        elif "scope" in scope:
+            scoped_parsed.append(scope.get("scope", ""))
+        else:
+            raise CLITokenException("Unknown error")
+
+    try:
+        await client.create_token(name, scoped_parsed, origin_code=None)
+    except AuthNoTokenException:
+        raise
+    except Exception as e:
+        raise CLITokenException(FeedbackManager.error_exception(error=e))
+
+    click.echo("The token has been generated successfully.")

tinybird/tb_cli_modules/workspace.py

@@ -0,0 +1,269 @@
+# This is a command file for our CLI. Please keep it clean.
+#
+# - If it makes sense and only when strictly necessary, you can create utility functions in this file.
+# - But please, **do not** interleave utility functions and command definitions.
+
+from typing import Optional
+
+import click
+from click import Context
+
+from tinybird.client import CanNotBeDeletedException, DoesNotExistException, TinyB
+from tinybird.config import get_display_host
+from tinybird.datafile import PipeTypes
+from tinybird.feedback_manager import FeedbackManager
+from tinybird.tb_cli_modules.cli import cli
+from tinybird.tb_cli_modules.common import (
+    _get_workspace_plan_name,
+    ask_for_user_token,
+    check_user_token,
+    coro,
+    create_workspace_interactive,
+    create_workspace_non_interactive,
+    echo_safe_humanfriendly_tables_format_smart_table,
+    get_current_main_workspace,
+    is_valid_starterkit,
+    print_current_workspace,
+    switch_workspace,
+)
+from tinybird.tb_cli_modules.config import CLIConfig
+from tinybird.tb_cli_modules.exceptions import CLIWorkspaceException
+
+
+@cli.group()
+@click.pass_context
+def workspace(ctx: Context) -> None:
+    """Workspace commands"""
+
+
+@workspace.command(name="ls")
+@click.pass_context
+@coro
+async def workspace_ls(ctx: Context) -> None:
+    """List all the workspaces you have access to in the account you're currently authenticated into."""
+
+    config = CLIConfig.get_project_config()
+    client = config.get_client()
+
+    response = await client.user_workspaces()
+
+    current_main_workspace = await get_current_main_workspace(config)
+    if not current_main_workspace:
+        raise CLIWorkspaceException(FeedbackManager.error_unable_to_identify_main_workspace())
+
+    columns = ["name", "id", "role", "plan", "current"]
+    table = []
+    click.echo(FeedbackManager.info_workspaces())
+
+    for workspace in response["workspaces"]:
+        table.append(
+            [
+                workspace["name"],
+                workspace["id"],
+                workspace["role"],
+                _get_workspace_plan_name(workspace["plan"]),
+                current_main_workspace["id"] == workspace["id"],
+            ]
+        )
+
+    echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
+
+
+@workspace.command(name="use")
+@click.argument("workspace_name_or_id")
+@coro
+async def workspace_use(workspace_name_or_id: str) -> None:
+    """Switch to another workspace. Use 'tb workspace ls' to list the workspaces you have access to."""
+    config = CLIConfig.get_project_config()
+
+    await switch_workspace(config, workspace_name_or_id)
+
+
+@workspace.command(name="current")
+@coro
+async def workspace_current():
+    """Show the workspace you're currently authenticated to"""
+    config = CLIConfig.get_project_config()
+
+    await print_current_workspace(config)
+
+
+@workspace.command(
+    name="clear",
+    short_help="Drop all the resources inside a project. This command is dangerous because it removes everything, use with care.",
+)
+@click.option("--yes", is_flag=True, default=False, help="Do not ask for confirmation")
+@click.option("--dry-run", is_flag=True, default=False, help="Run the command without removing anything")
+@click.pass_context
+@coro
+async def clear_workspace(ctx: Context, yes: bool, dry_run: bool) -> None:
+    """Drop all the resources inside a project. This command is dangerous because it removes everything, use with care."""
+
+    # Get current workspace to add the name to the alert message
+    config = CLIConfig.get_project_config()
+    client: TinyB = config.get_client()
+
+    response = await client.user_workspaces_and_branches()
+
+    columns = ["name", "id", "role", "plan", "current"]
+    table = []
+
+    for workspace in response["workspaces"]:
+        if config["id"] == workspace["id"]:
+            if workspace.get("is_branch"):
+                raise CLIWorkspaceException(FeedbackManager.error_not_allowed_in_branch())
+                return
+            else:
+                click.echo(FeedbackManager.info_current_workspace())
+                table.append(
+                    [
+                        workspace["name"],
+                        workspace["id"],
+                        workspace["role"],
+                        _get_workspace_plan_name(workspace["plan"]),
+                        True,
+                    ]
+                )
+            break
+
+    echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
+
+    if yes or click.confirm(FeedbackManager.warning_confirm_clear_workspace()):
+        pipes = await client.pipes(dependencies=False, node_attrs="id,name,materialized", attrs="name,type")
+        pipe_names = [pipe["name"] for pipe in pipes]
+
+        for pipe in pipes:
+            if pipe["type"] == PipeTypes.MATERIALIZED:
+                if not dry_run:
+                    node_id = None
+                    for node in pipe["nodes"]:
+                        if "materialized" in node and node["materialized"] is not None:
+                            node_id = node["id"]
+                            break
+
+                    if node_id:
+                        click.echo(FeedbackManager.info_unlinking_materialized_pipe(pipe=pipe["name"]))
+                        try:
+                            await client.pipe_unlink_materialized(pipe["name"], node_id)
+                        except DoesNotExistException:
+                            click.echo(FeedbackManager.info_materialized_unlinking_pipe_not_found(pipe=pipe["name"]))
+                else:
+                    click.echo(FeedbackManager.info_materialized_dry_unlinking_pipe(pipe=pipe["name"]))
+
+        for pipe_name in pipe_names:
+            if not dry_run:
+                click.echo(FeedbackManager.info_removing_pipe(pipe=pipe_name))
+                try:
+                    await client.pipe_delete(pipe_name)
+                except DoesNotExistException:
+                    click.echo(FeedbackManager.info_removing_pipe_not_found(pipe=pipe_name))
+            else:
+                click.echo(FeedbackManager.info_dry_removing_pipe(pipe=pipe_name))
+
+        datasources = await client.datasources()
+        ds_names = [datasource["name"] for datasource in datasources]
+        for ds_name in ds_names:
+            if not dry_run:
+                click.echo(FeedbackManager.info_removing_datasource(datasource=ds_name))
+                try:
+                    await client.datasource_delete(ds_name, force=True)
+                except DoesNotExistException:
+                    click.echo(FeedbackManager.info_removing_datasource_not_found(datasource=ds_name))
+                except CanNotBeDeletedException as e:
+                    raise CLIWorkspaceException(
+                        FeedbackManager.error_datasource_can_not_be_deleted(datasource=ds_name, error=e)
+                    )
+                except Exception as e:
+                    if "is a Shared Data Source" in str(e):
+                        raise CLIWorkspaceException(FeedbackManager.error_operation_can_not_be_performed(error=e))
+                    else:
+                        raise CLIWorkspaceException(FeedbackManager.error_exception(error=e))
+            else:
+                click.echo(FeedbackManager.info_dry_removing_datasource(datasource=ds_name))
+
+
+@workspace.command(name="create", short_help="Create a new Workspace for your Tinybird user")
+@click.argument("workspace_name", required=False)
+@click.option("--starter_kit", "starter_kit", type=str, required=False, help="Use a Tinybird starter kit as a template")
+@click.option("--starter-kit", "starter_kit", hidden=True)
+@click.option("--user_token", is_flag=False, default=None, help="When passed, we won't prompt asking for it")
+@click.option(
+    "--fork",
+    is_flag=True,
+    default=False,
+    help="When enabled, we will share all datasource from the current workspace to the new created one",
+)
+@click.pass_context
+@coro
+async def create_workspace(
+    ctx: Context, workspace_name: str, starter_kit: str, user_token: Optional[str], fork: bool
+) -> None:
+    if starter_kit and not await is_valid_starterkit(ctx, starter_kit):
+        raise CLIWorkspaceException(FeedbackManager.error_starterkit_name(starterkit_name=starter_kit))
+
+    if not user_token:
+        config = CLIConfig.get_project_config()
+        host = config.get_host() or CLIConfig.DEFAULTS["host"]
+        ui_host = get_display_host(host)
+        user_token = ask_for_user_token("create a new workspace", ui_host)
+        if not user_token:
+            return
+    await check_user_token(ctx, user_token)
+
+    # If we have at least workspace_name, we start the non interactive
+    # process, creating an empty workspace
+    if workspace_name:
+        await create_workspace_non_interactive(ctx, workspace_name, starter_kit, user_token, fork)
+    else:
+        await create_workspace_interactive(ctx, workspace_name, starter_kit, user_token, fork)
+
+
+@workspace.command(name="delete", short_help="Delete a Workspace for your Tinybird user")
+@click.argument("workspace_name_or_id")
+@click.option("--user_token", is_flag=False, default=None, help="When passed, we won't prompt asking for it")
+@click.option(
+    "--confirm_hard_delete",
+    default=None,
+    help="Introduce the name of the workspace in order to confirm you want to run a hard delete over the workspace",
+    hidden=True,
+)
+@click.option("--yes", is_flag=True, default=False, help="Do not ask for confirmation")
+@click.pass_context
+@coro
+async def delete_workspace(
+    ctx: Context, workspace_name_or_id: str, user_token: Optional[str], confirm_hard_delete: Optional[str], yes: bool
+) -> None:
+    """Delete a workspace where you are an admin."""
+
+    config = CLIConfig.get_project_config()
+    client = config.get_client()
+    host = config.get_host() or CLIConfig.DEFAULTS["host"]
+    ui_host = get_display_host(host)
+
+    if not user_token:
+        user_token = ask_for_user_token("delete a workspace", ui_host)
+    await check_user_token(ctx, user_token)
+
+    workspaces = (await client.user_workspaces()).get("workspaces", [])
+    workspace_to_delete = next(
+        (
+            workspace
+            for workspace in workspaces
+            if workspace["name"] == workspace_name_or_id or workspace["id"] == workspace_name_or_id
+        ),
+        None,
+    )
+
+    if not workspace_to_delete:
+        raise CLIWorkspaceException(FeedbackManager.error_workspace(workspace=workspace_name_or_id))
+
+    if yes or click.confirm(
+        FeedbackManager.warning_confirm_delete_workspace(workspace_name=workspace_to_delete.get("name"))
+    ):
+        client.token = user_token
+
+        try:
+            await client.delete_workspace(workspace_to_delete["id"], confirm_hard_delete)
+            click.echo(FeedbackManager.success_workspace_deleted(workspace_name=workspace_to_delete["name"]))
+        except Exception as e:
+            raise CLIWorkspaceException(FeedbackManager.error_exception(error=str(e)))