Example #1
0
def load_flows_from_module(name: str) -> "List[prefect.Flow]":
    """
    Given a module name (or full import path to a flow), load all flows found in the
    module
    """
    # TODO: This is copied and slightly modified from `prefect.cli.build_register`
    #       we should probably abstract this in the future
    try:
        with prefect.context({"loading_flow": True}):
            mod_or_obj = import_object(name)
    except Exception as exc:
        # If the requested module isn't found, log without a traceback
        # otherwise log a general message with the traceback.
        if isinstance(exc, ModuleNotFoundError) and (
                name == exc.name or
            (name.startswith(exc.name) and name[len(exc.name)] == ".")):
            raise TerminalError(str(exc).capitalize())
        elif isinstance(exc, AttributeError):
            raise TerminalError(str(exc).capitalize())
        else:
            raise

    if isinstance(mod_or_obj, ModuleType):
        flows = [
            f for f in vars(mod_or_obj).values()
            if isinstance(f, prefect.Flow)
        ]
    elif isinstance(mod_or_obj, prefect.Flow):
        flows = [mod_or_obj]
    else:
        raise TerminalError(
            f"Invalid object of type {type(mod_or_obj).__name__!r} found at {name!r}. "
            f"Expected Module or Flow.")

    return flows
Example #2
0
def get_flow_from_path_or_module(path: str = None,
                                 module: str = None,
                                 name: str = None):
    location = path if path is not None else module
    flows = load_flows_from_script(path) if path else load_flows_from_module(
        module)
    flows_by_name = {flow.name: flow for flow in flows}
    flow_names = ", ".join(map(repr, flows_by_name.keys()))

    if not flows:
        raise TerminalError(f"Found no flows at {location}.")

    if len(flows) > 1 and not name:
        raise TerminalError(
            f"Found multiple flows at {location!r}: {flow_names}\n\n"
            f"Specify a flow name to run.")
    if name:
        if name not in flows_by_name:
            raise TerminalError(
                f"Did not find {name!r} in flows at {location}. Found {flow_names}"
            )
        flow = flows_by_name[name]

    else:
        flow = list(flows_by_name.values())[0]

    return flow
Example #3
0
def load_json_key_values(
        cli_input: List[str],
        display_name: str) -> Dict[str, Union[dict, str, int]]:
    """
    Parse a list of strings formatted as "key=value" where the value is loaded as JSON.

    We do the best here to display a helpful JSON parsing message, e.g.
    ```
    Error: Failed to parse JSON for parameter 'name' with value

        foo

    JSON Error: Expecting value: line 1 column 1 (char 0)
    Did you forget to include quotes? You may need to escape so your shell does not remove them, e.g. \"
    ```

    Args:
        cli_input: A list of "key=value" strings to parse
        display_name: A name to display in exceptions

    Returns:
        A mapping of keys -> parsed values
    """
    parsed = {}

    def cast_value(value: str) -> Any:
        """Cast the value from a string to a valid JSON type; add quotes for the user
        if necessary
        """
        try:
            return json.loads(value)
        except ValueError as exc:
            if ("Extra data" in str(exc)
                    or "Expecting value" in str(exc)) and '"' not in value:
                return cast_value(f'"{value}"')
            raise exc

    for spec in cli_input:
        try:
            key, value = spec.split("=")
        except ValueError:
            raise TerminalError(
                f"Invalid {display_name} option {spec!r}. Expected format 'key=value'."
            )

        try:
            parsed[key] = cast_value(value)
        except ValueError as exc:
            indented_value = textwrap.indent(value, prefix="\t")
            raise TerminalError(
                f"Failed to parse JSON for {display_name} {key!r} with value"
                f"\n\n{indented_value}\n\n"
                f"JSON Error: {exc}")

    return parsed
Example #4
0
def switch_tenants(id, slug, default):
    """
    Switch active tenant

    \b
    Options:
        --id, -i    TEXT    A Prefect Cloud tenant id
        --slug, -s  TEXT    A Prefect Cloud tenant slug
    """

    # If the config specifies a tenant explicitly, it is used before this mechanism
    if config.cloud.get("tenant_id"):
        raise TerminalError(
            "Your tenant id has been set in the Prefect config instead of with the "
            "CLI. To switch tenants with the CLI, remove the config key "
            " `prefect.cloud.tenant_id`")

    client = Client()

    # Deprecated API token check
    if not client.api_key:
        check_override_auth_token()

        if default:
            raise TerminalError(
                "The default tenant flag can only be used with API keys.")

    else:  # Using an API key
        if default:
            # Clear the set tenant on disk
            client.tenant_id = None
            client.save_auth_to_disk()
            click.secho(
                "Tenant restored to the default tenant for your API key: "
                f"{client.get_default_tenant()}",
                fg="green",
            )
            return

    login_success = client.login_to_tenant(tenant_slug=slug, tenant_id=id)
    if not login_success:
        raise TerminalError("Unable to switch tenant!")

    # `login_to_tenant` will write to disk if using an API token, if using an API key
    # we will write to disk manually here
    if client.api_key:
        client.save_auth_to_disk()

    click.secho(f"Tenant switched to {client.tenant_id}", fg="green")
Example #5
0
def abort_on_config_api_key(message: str = None):
    if config.cloud.get("api_key"):
        # Add a leading space if not null
        message = (" " + message) if message else ""
        raise TerminalError(
            "Your API key is set in the Prefect config instead of with the CLI."
            + message)
Example #6
0
def logout(token):
    """
    Log out of Prefect Cloud

    This will remove your cached authentication from disk.
    """

    client = Client()

    # Log out of API keys unless given the token flag
    if client.api_key and not token:

        # Check the source of the API key
        abort_on_config_api_key(
            "To log out, remove the config key `prefect.cloud.api_key`")

        click.confirm(
            "Are you sure you want to log out of Prefect Cloud? "
            "This will remove your API key from this machine.",
            default=False,
            abort=True,
        )

        # Clear the key and tenant id then write to the cache
        client.api_key = ""
        client._tenant_id = ""
        client.save_auth_to_disk()

        click.secho("Logged out of Prefect Cloud", fg="green")

    else:
        raise TerminalError("You are not logged in to Prefect Cloud. "
                            "Use `prefect auth login` to log in first.")
Example #7
0
def list_keys():
    """
    List available Prefect Cloud API keys.
    """
    client = Client()

    response = client.graphql(
        query={
            "query": {
                "auth_api_key": {
                    "id": True,
                    "name": True,
                    "expires_at": True,
                }
            }
        })
    keys = response.get("data", {}).get("auth_api_key")
    if keys is None:
        raise TerminalError(
            f"Unexpected response from Prefect Cloud: {response}")

    if not keys:
        click.secho("You have not created any API keys", fg="yellow")

    else:
        click.echo(
            tabulate(
                [(key.name, key.id, key.expires_at or "NEVER")
                 for key in keys],
                headers=["NAME", "ID", "EXPIRES AT"],
                tablefmt="plain",
                numalign="left",
                stralign="left",
            ))
Example #8
0
def list_tenants():
    """
    List available tenants
    """
    client = Client()

    try:
        tenants = client.get_available_tenants()
    except AuthorizationError:
        raise TerminalError(
            "You are not authenticated. Use `prefect auth login` first.")

    output = []
    for item in tenants:
        active = None
        if item.id == client.tenant_id:
            active = "*"
        output.append([item.name, item.slug, item.id, active])

    click.echo(
        tabulate(
            output,
            headers=["NAME", "SLUG", "ID", ""],
            tablefmt="plain",
            numalign="left",
            stralign="left",
        ))
Example #9
0
def load_flows_from_script(path: str) -> "List[prefect.Flow]":
    """Given a file path, load all flows found in the file"""
    # TODO: This is copied and slightly modified from `prefect.cli.build_register`
    #       we should probably abstract this in the future
    # Temporarily add the flow's local directory to `sys.path` so that local
    # imports work. This ensures that `sys.path` is the same as it would be if
    # the flow script was run directly (i.e. `python path/to/flow.py`).
    orig_sys_path = sys.path.copy()
    sys.path.insert(0, os.path.dirname(os.path.abspath(path)))
    try:
        with prefect.context({
                "loading_flow": True,
                "local_script_path": path
        }):
            namespace = runpy.run_path(path, run_name="<flow>")
    except FileNotFoundError as exc:
        if path in str(exc):  # Only capture it if it's about our file
            raise TerminalError(
                f"File does not exist: {os.path.abspath(path)!r}")
        raise
    finally:
        sys.path[:] = orig_sys_path

    flows = [f for f in namespace.values() if isinstance(f, prefect.Flow)]
    return flows
Example #10
0
def login(key):
    """
    Login to Prefect Cloud

    Create an API key in the UI then login with it here:

        $ prefect auth login -k YOUR-KEY

    You will be switched to the default tenant associated with the key. After login,
    your available tenants can be seen with `prefect auth list-tenants` and you can
    change the default tenant on this machine using `prefect auth switch-tenants`.

    The given key will be stored on disk for later access. Prefect will default to using
    this key for all interaction with the API but frequently overrides can be passed to
    individual commands or functions. To remove your key from disk, see
    `prefect auth logout`.
    """
    if not key:
        raise TerminalError("You must supply an API key!")

    abort_on_config_api_key(
        "To log in with the CLI, remove the config key `prefect.cloud.api_key`"
    )

    # Ignore any tenant id that has been previously set via login
    client = Client(api_key=key)
    client._tenant_id = None

    try:
        tenant_id = client._get_auth_tenant()
    except AuthorizationError:
        raise TerminalError("Unauthorized. Invalid Prefect Cloud API key.")
    except ClientError:
        raise TerminalError(
            "Error attempting to communicate with Prefect Cloud.")
    else:
        client.tenant_id = tenant_id
        client.save_auth_to_disk()
        tenant = TenantView.from_tenant_id(tenant_id)
        click.secho(
            f"Logged in to Prefect Cloud tenant {tenant.name!r} ({tenant.slug})",
            fg="green",
        )
        return
Example #11
0
def switch_tenants(id, slug, default):
    """
    Switch active tenant

    \b
    Options:
        --id, -i    TEXT    A Prefect Cloud tenant id
        --slug, -s  TEXT    A Prefect Cloud tenant slug
    """

    # If the config specifies a tenant explicitly, it is used before this mechanism
    if config.cloud.get("tenant_id"):
        raise TerminalError(
            "Your tenant id has been set in the Prefect config instead of with the "
            "CLI. To switch tenants with the CLI, remove the config key "
            " `prefect.cloud.tenant_id`")

    client = Client()

    if not client.api_key:
        raise TerminalError("You are not logged in!")

    if default:
        # Clear the set tenant on disk
        client.tenant_id = None
        client.save_auth_to_disk()
        click.secho(
            "Tenant restored to the default tenant for your API key: "
            f"{client._get_auth_tenant()}",
            fg="green",
        )
        return

    try:
        tenant_id = client.switch_tenant(tenant_slug=slug, tenant_id=id)
    except AuthorizationError:
        raise TerminalError(
            "Unauthorized. Your API key is not valid for that tenant.")

    client.save_auth_to_disk()
    click.secho(f"Tenant switched to {tenant_id}", fg="green")
Example #12
0
def list_command():
    """
    List all key value pairs
    """
    try:
        result = kv_store.list_keys()
        if result:
            click.secho("\n".join(result), fg="green")
        else:
            click.secho("No keys found", fg="yellow")
    except Exception as exc:
        log_exception(exc)
        raise TerminalError("An error occurred when listing keys")
Example #13
0
def get_flow_view(
    flow_or_group_id: str = None,
    project: str = None,
    name: str = None,
) -> "FlowView":
    if flow_or_group_id:
        # Lookup by flow id then flow group id if that fails
        try:
            flow_view = FlowView.from_flow_id(flow_or_group_id)
        except ValueError:
            pass
        else:
            return flow_view

        try:
            flow_view = FlowView.from_flow_group_id(flow_or_group_id)
        except ValueError:
            pass
        else:
            return flow_view

        # Fall through to failure
        raise TerminalError(
            f"Failed to find flow id or flow group id matching {flow_or_group_id!r}"
        )

    if project:
        if not name:
            raise TerminalError(
                "Missing required option `--name`. Cannot look up a flow by project "
                "without also passing a name.")
        return FlowView.from_flow_name(flow_name=name, project_name=project)

    if name:
        # If name wasn't provided for use with another lookup, try a global name search
        return FlowView.from_flow_name(flow_name=name)

    # This line should not be reached
    raise RuntimeError("Failed to find matching case for flow lookup.")
Example #14
0
def delete_command(key):
    """
    Delete a key value pair

    \b
    Arguments:
        key         TEXT    Key to delete
    """
    try:
        kv_store.delete_key(key=key)
        click.secho(f"Key {key!r} has been deleted", fg="green")
    except Exception as exc:
        log_exception(exc)
        raise TerminalError("An error occurred deleting the key")
Example #15
0
def get_command(key):
    """
    Get the value of a key

    \b
    Arguments:
        key         TEXT    Key to get
    """
    try:
        result = kv_store.get_key_value(key=key)
        click.secho(f"Key {key!r} has value {result!r}", fg="green")
    except Exception as exc:
        log_exception(exc)
        raise TerminalError(f"Error retrieving value for key {key!r}")
Example #16
0
def set_command(key, value):
    """
    Set a key value pair, overriding existing values if key exists

    \b
    Arguments:
        key         TEXT    Key to set
        value       TEXT    Value associated with key to set
    """
    try:
        kv_store.set_key_value(key=key, value=value)
        click.secho("Key value pair set successfully", fg="green")
    except Exception as exc:
        log_exception(exc)
        raise TerminalError("An error occurred setting the key value pair")
Example #17
0
def revoke_key(id):
    """
    Revoke a Prefect Cloud API key.
    """
    client = Client()

    output = client.graphql(
        query={
            "mutation($input: delete_api_key_input!)": {
                "delete_api_key(input: $input)": {"success"}
            }
        },
        variables=dict(input=dict(key_id=id)),
    )

    if not output.get("data", None) or not output.data.delete_api_key.success:
        raise TerminalError(f"Unable to revoke key {id!r}")

    click.secho("Key successfully revoked!", fg="green")
Example #18
0
def try_error_done(
    message: str,
    echo: Callable = click.secho,
    traceback: bool = False,
    skip_done: bool = False,
):
    """
    Try to run the code in the context block. On error print "Error" and raise a
    terminal error with the exception string. On succecss, print "Done".

    Args:
        message: The first message to display
        echo: The function to use to echo. Must support `click.secho` arguments
        traceback: Display the exception traceback instead of a short message
        skip_done: Do not display 'Done', the user of the context should instead

    Example:
        >>> with try_error_done("Setting up foo..."):
        >>>    pass
        Setting up foo... Done
        >>> with try_error_done("Setting up bar..."):
        >>>    raise ValueError("no!")
        Setting up bar... Error
        no!
    """
    echo(message, nl=False)
    try:
        yield
    except TerminalError:
        echo(" Error", fg="red")
        raise
    except Exception as exc:
        echo(" Error", fg="red")

        if traceback and not isinstance(exc, TerminalError):
            log_exception(exc, indent=2)
            raise TerminalError
        else:
            raise TerminalError(f"{type(exc).__name__}: {exc}")

    else:
        if not skip_done:
            echo(" Done", fg="green")
Example #19
0
def login(key, token):
    """
    Login to Prefect Cloud

    Create an API key in the UI then login with it here:

        $ prefect auth login -k YOUR-KEY

    You will be switched to the default tenant associated with the key. After login,
    your available tenants can be seen with `prefect auth list-tenants` and you can
    change the default tenant on this machine using `prefect auth switch-tenants`.

    The given key will be stored on disk for later access. Prefect will default to using
    this key for all interaction with the API but frequently overrides can be passed to
    individual commands or functions. To remove your key from disk, see
    `prefect auth logout`.

    This command has backwards compatibility support for API tokens, which are a
    deprecated form of authentication with Prefect Cloud
    """
    if not key and not token:
        raise TerminalError("You must supply an API key or token!")

    if key and token:
        raise TerminalError("You cannot supply both an API key and token")

    abort_on_config_api_key(
        "To log in with the CLI, remove the config key `prefect.cloud.api_key`"
    )

    # Attempt to treat the input like an API key even if it is passed as a token
    client = Client(api_key=key or token)

    try:
        default_tenant = client.get_default_tenant()
    except AuthorizationError:
        if key:  # We'll catch an error again later if using a token
            raise TerminalError("Unauthorized. Invalid Prefect Cloud API key.")
    except ClientError:
        raise TerminalError(
            "Error attempting to communicate with Prefect Cloud.")
    else:
        if not default_tenant and key:
            raise TerminalError(
                "Failed to find a tenant associated with the given API key!")

        elif default_tenant:  # Successful login
            if token:
                click.secho(
                    "WARNING: You logged in with an API key using the `--token` flag "
                    "which is deprecated. Please use `--key` instead.",
                    fg="yellow",
                )
            client.save_auth_to_disk()
            click.secho("Login successful!", fg="green")
            return

        # If there's not a tenant id, we've been given an actual token, fallthrough to
        # the backwards compatibility token auth

    # Backwards compatibility for tokens
    if token:
        check_override_auth_token()
        client = Client(api_token=token)

        # Verify they're not also using an API key
        if client.api_key:
            raise TerminalError(
                "You have already logged in with an API key and cannot use a token."
            )

        click.secho(
            "WARNING: API tokens are deprecated. Please create an API key and use "
            "`prefect auth login --key <KEY>` to login instead.",
            fg="yellow",
        )

        # Verify login obtained a valid api token
        try:
            output = client.graphql(
                query={"query": {
                    "user": {
                        "default_membership": "tenant_id"
                    }
                }})

            # Log into default membership
            success_login = client.login_to_tenant(
                tenant_id=output.data.user[0].default_membership.tenant_id)

            if not success_login:
                raise AuthorizationError

        except AuthorizationError:
            click.secho(
                "Error attempting to use the given API token. "
                "Please check that you are providing a USER scoped Personal Access Token.\n"
                "For more information visit the documentation for USER tokens at "
                "https://docs.prefect.io/orchestration/concepts/tokens.html#user",
                fg="red",
            )
            return
        except ClientError:
            click.secho(
                "Error attempting to communicate with Prefect Cloud. "
                "Please check that you are providing a USER scoped Personal Access Token.\n"
                "For more information visit the documentation for USER tokens at "
                "https://docs.prefect.io/orchestration/concepts/tokens.html#user",
                fg="red",
            )
            return

        # save token
        client.save_api_token()

        click.secho("Login successful!", fg="green")
Example #20
0
def create_key(name, expire, quiet):
    """
    Create a Prefect Cloud API key for authentication with your current user
    """
    # TODO: Add service account associated key creation eventually

    # Parse the input expiration
    if expire is not None:
        try:
            expires_at = pendulum.parse(expire, strict=False)
        except pendulum.parsing.exceptions.ParserError as exc:
            raise TerminalError(
                f"Failed to parse expiration time. {exc}\n"
                "Please pass a date in a dateutil parsable format.")

        if expires_at.diff(abs=False).in_seconds() > 0:
            raise TerminalError(
                f"Given expiration time {expire!r} is a time in the past: {expires_at}"
            )
        expire_msg = f" that will expire {expires_at.diff_for_humans()}"
    else:
        expires_at = None
        expire_msg = ""

    client = Client()

    # We must retrieve our own user id first since you could be creating a key for a SA
    if not quiet:
        click.echo("Retrieving user information...")

    response = client.graphql({"query": {"auth_info": {"user_id"}}})
    user_id = response.get("data", {}).get("auth_info", {}).get("user_id")
    if not user_id:
        raise TerminalError(
            "Failed to retrieve the current user id from Prefect Cloud")

    # Actually create the key
    if not quiet:
        click.echo(f"Creating key{expire_msg}...")
    response = client.graphql(
        query={
            "mutation($input: create_api_key_input!)": {
                "create_api_key(input: $input)": {"key"}
            }
        },
        variables=dict(input=dict(
            name=name,
            user_id=user_id,
            expires_at=expires_at.in_tz("utc").isoformat(
            ) if expires_at else None,
        )),
    )

    key = response.get("data", {}).get("create_api_key", {}).get("key")
    if key is None:
        raise TerminalError(
            f"Unexpected response from Prefect Cloud: {response}")

    if quiet:
        click.echo(key)
    else:
        click.echo(
            "This is the only time this key will be displayed! Store it somewhere safe."
        )
        click.secho(f"Successfully created key: {key}", fg="green")
Example #21
0
def run(
    ctx,
    flow_or_group_id,
    project,
    path,
    module,
    name,
    labels,
    context_vars,
    params,
    execute,
    idempotency_key,
    schedule,
    log_level,
    param_file,
    run_name,
    quiet,
    no_logs,
    watch,
):
    """Run a flow"""
    # Since the old command was a subcommand of this, we have to do some
    # mucking to smoothly deprecate it. Can be removed with `prefect run flow`
    # is removed.
    if ctx.invoked_subcommand is not None:
        if any([params, no_logs, quiet, flow_or_group_id]):
            # These options are not supported by `prefect run flow`
            raise ClickException("Got unexpected extra argument (%s)" %
                                 ctx.invoked_subcommand)
        return

    # Define a simple function so we don't have to have a lot of `if not quiet` logic
    quiet_echo = ((lambda *_, **__: None) if quiet else
                  lambda *args, **kwargs: click.secho(*args, **kwargs))

    # Cast labels to a list instead of a tuple so we can extend it
    labels = list(labels)

    # Ensure that the user has not passed conflicting options
    given_lookup_options = {
        key
        for key, option in {
            "--id": flow_or_group_id,
            "--project": project,
            "--path": path,
            "--module": module,
        }.items() if option is not None
    }
    # Since `name` can be passed in conjunction with several options and also alone
    # it requires a special case here
    if not given_lookup_options and not name:
        raise ClickException("Received no options to look up the flow." +
                             FLOW_LOOKUP_MSG)
    if "--id" in given_lookup_options and name:
        raise ClickException("Received too many options to look up the flow; "
                             "cannot specifiy both `--name` and `--id`" +
                             FLOW_LOOKUP_MSG)
    if len(given_lookup_options) > 1:
        raise ClickException("Received too many options to look up the flow: "
                             f"{', '.join(given_lookup_options)}" +
                             FLOW_LOOKUP_MSG)

    # Load parameters and context ------------------------------------------------------
    context_dict = load_json_key_values(context_vars, "context")

    file_params = {}
    if param_file:

        try:
            with open(param_file) as fp:
                file_params = json.load(fp)
        except FileNotFoundError:
            raise TerminalError(
                f"Parameter file does not exist: {os.path.abspath(param_file)!r}"
            )
        except ValueError as exc:
            raise TerminalError(
                f"Failed to parse JSON at {os.path.abspath(param_file)!r}: {exc}"
            )

    cli_params = load_json_key_values(params, "parameter")
    conflicting_keys = set(cli_params.keys()).intersection(file_params.keys())
    if conflicting_keys:
        quiet_echo(
            "The following parameters were specified by file and CLI, the CLI value "
            f"will be used: {conflicting_keys}")
    params_dict = {**file_params, **cli_params}

    # Local flow run -------------------------------------------------------------------

    if path or module:
        # We can load a flow for local execution immediately if given a path or module,
        # otherwise, we'll lookup the flow then pull from storage for a local run
        with try_error_done("Retrieving local flow...",
                            quiet_echo,
                            traceback=True):
            flow = get_flow_from_path_or_module(path=path,
                                                module=module,
                                                name=name)

        # Set the desired log level
        if no_logs:
            log_level = 100  # CRITICAL is 50 so this should do it

        run_info = ""
        if params_dict:
            run_info += f"└── Parameters: {params_dict}\n"
        if context_dict:
            run_info += f"└── Context: {context_dict}\n"

        if run_info:
            quiet_echo("Configured local flow run")
            quiet_echo(run_info, nl=False)

        quiet_echo("Running flow locally...")
        with temporary_logger_config(
                level=log_level,
                stream_fmt="└── %(asctime)s | %(levelname)-7s | %(message)s",
                stream_datefmt="%H:%M:%S",
        ):
            with prefect.context(**context_dict):
                try:
                    result_state = flow.run(parameters=params_dict,
                                            run_on_schedule=schedule)
                except Exception as exc:
                    quiet_echo("Flow runner encountered an exception!")
                    log_exception(exc, indent=2)
                    raise TerminalError("Flow run failed!")

        if result_state.is_failed():
            quiet_echo("Flow run failed!", fg="red")
            sys.exit(1)
        else:
            quiet_echo("Flow run succeeded!", fg="green")

        return

    # Backend flow run -----------------------------------------------------------------

    if schedule:
        raise ClickException(
            "`--schedule` can only be specified for local flow runs")

    client = Client()

    # Validate the flow look up options we've been given and get the flow from the
    # backend
    with try_error_done("Looking up flow metadata...", quiet_echo):
        flow_view = get_flow_view(
            flow_or_group_id=flow_or_group_id,
            project=project,
            name=name,
        )

    if log_level:
        run_config = flow_view.run_config
        if not run_config.env:
            run_config.env = {}
        run_config.env["PREFECT__LOGGING__LEVEL"] = log_level
    else:
        run_config = None

    if execute:
        # Add a random label to prevent an agent from picking up this run
        labels.append(f"agentless-run-{str(uuid.uuid4())[:8]}")

    try:  # Handle keyboard interrupts during creation
        flow_run_id = None

        # Create a flow run in the backend
        with try_error_done(
                f"Creating run for flow {flow_view.name!r}...",
                quiet_echo,
                traceback=True,
                # Display 'Done' manually after querying for data to display so there is not
                # a lag
                skip_done=True,
        ):
            flow_run_id = client.create_flow_run(
                flow_id=flow_view.flow_id,
                parameters=params_dict,
                context=context_dict,
                # If labels is an empty list pass `None` to get defaults
                # https://github.com/PrefectHQ/server/blob/77c301ce0c8deda4f8771f7e9991b25e7911224a/src/prefect_server/api/runs.py#L136
                labels=labels or None,
                run_name=run_name,
                # We only use the run config for setting logging levels right now
                run_config=run_config,
                idempotency_key=idempotency_key,
            )

        if quiet:
            # Just display the flow run id in quiet mode
            click.echo(flow_run_id)
            flow_run = None
        else:
            # Grab information about the flow run (if quiet we can skip this query)
            flow_run = FlowRunView.from_flow_run_id(flow_run_id)
            run_url = client.get_cloud_url("flow-run", flow_run_id)

            # Display "Done" for creating flow run after pulling the info so there
            # isn't a weird lag
            quiet_echo(" Done", fg="green")
            quiet_echo(
                textwrap.dedent(f"""
                        └── Name: {flow_run.name}
                        └── UUID: {flow_run.flow_run_id}
                        └── Labels: {flow_run.labels}
                        └── Parameters: {flow_run.parameters}
                        └── Context: {flow_run.context}
                        └── URL: {run_url}
                        """).strip())

    except KeyboardInterrupt:
        # If the user interrupts here, they will expect the flow run to be cancelled
        quiet_echo("\nKeyboard interrupt detected! Aborting...", fg="yellow")
        if flow_run_id:
            client.cancel_flow_run(flow_run_id=flow_run_id)
            quiet_echo("Cancelled flow run.")
        else:
            # The flow run was not created so we can just exit
            quiet_echo("Aborted.")
        return

    # Handle agentless execution
    if execute:
        quiet_echo("Executing flow run...")
        try:
            with temporary_logger_config(
                    level=(100 if no_logs or quiet else
                           log_level),  # Disable logging if asked
                    stream_fmt=
                    "└── %(asctime)s | %(levelname)-7s | %(message)s",
                    stream_datefmt="%H:%M:%S",
            ):
                execute_flow_run_in_subprocess(flow_run_id)
        except KeyboardInterrupt:
            quiet_echo("Keyboard interrupt detected! Aborting...", fg="yellow")
            pass

    elif watch:
        try:
            quiet_echo("Watching flow run execution...")
            for log in watch_flow_run(
                    flow_run_id=flow_run_id,
                    stream_logs=not no_logs,
            ):
                level_name = logging.getLevelName(log.level)
                timestamp = log.timestamp.in_tz(tz="local")
                echo_with_log_color(
                    log.level,
                    f"└── {timestamp:%H:%M:%S} | {level_name:<7} | {log.message}",
                )

        except KeyboardInterrupt:
            quiet_echo("Keyboard interrupt detected!", fg="yellow")
            try:
                cancel = click.confirm(
                    "On exit, we can leave your flow run executing or cancel it.\n"
                    "Do you want to cancel this flow run?",
                    default=True,
                )
            except click.Abort:
                # A second keyboard interrupt will exit without cancellation
                pass
            else:
                if cancel:
                    client.cancel_flow_run(flow_run_id=flow_run_id)
                    quiet_echo("Cancelled flow run.", fg="green")
                    return

            quiet_echo("Exiting without cancelling flow run!", fg="yellow")
            raise  # Re-raise the interrupt

    else:
        # If not watching or executing, exit without checking state
        return

    # Get the final flow run state
    flow_run = FlowRunView.from_flow_run_id(flow_run_id)

    # Wait for the flow run to be done up to 3 seconds
    elapsed_time = 0
    while not flow_run.state.is_finished() and elapsed_time < 3:
        time.sleep(1)
        elapsed_time += 1
        flow_run = flow_run.get_latest()

    # Display the final state
    if flow_run.state.is_failed():
        quiet_echo("Flow run failed!", fg="red")
        sys.exit(1)
    elif flow_run.state.is_successful():
        quiet_echo("Flow run succeeded!", fg="green")
    else:
        quiet_echo(f"Flow run is in unexpected state: {flow_run.state}",
                   fg="yellow")
        sys.exit(1)
Example #22
0
def logout(token):
    """
    Log out of Prefect Cloud

    This will remove your cached authentication from disk.
    """

    client = Client()

    # Log out of API keys unless given the token flag
    if client.api_key and not token:

        # Check the source of the API key
        abort_on_config_api_key(
            "To log out, remove the config key `prefect.cloud.api_key`")

        click.confirm(
            "Are you sure you want to log out of Prefect Cloud? "
            "This will remove your API key from this machine.",
            default=False,
            abort=True,
        )

        # Clear the key and tenant id then write to the cache
        client.api_key = ""
        client._tenant_id = ""
        client.save_auth_to_disk()

        click.secho("Logged out of Prefect Cloud", fg="green")

    elif client._api_token:

        check_override_auth_token()
        tenant_id = client.active_tenant_id

        if not tenant_id:
            click.confirm(
                "Are you sure you want to log out of Prefect Cloud? "
                "This will remove your API token from this machine.",
                default=False,
                abort=True,
            )

            # Remove the token from local storage by writing blank settings
            client._save_local_settings({})
            click.secho("Logged out of Prefect Cloud", fg="green")

        else:
            # Log out of the current tenant (dropping the access token) while retaining
            # the API token. This is backwards compatible behavior. Running the logout
            # command twice will remove the token from storage entirely
            click.confirm(
                "Are you sure you want to log out of your current Prefect Cloud tenant?",
                default=False,
                abort=True,
            )

            client.logout_from_tenant()

            click.secho(
                f"Logged out from tenant {tenant_id}. Run `prefect auth logout` again "
                "to delete your API token.",
                fg="green",
            )

    else:
        raise TerminalError("You are not logged in to Prefect Cloud. "
                            "Use `prefect auth login` to log in first.")