dvc/dvc/commands/studio.py

import os

from funcy import get_in

from dvc.cli import formatter
from dvc.cli.utils import append_doc_link
from dvc.commands.config import CmdConfig
from dvc.exceptions import DvcException
from dvc.log import logger

logger = logger.getChild(__name__)


class CmdStudioLogin(CmdConfig):
    def run(self):
        from dvc.env import DVC_STUDIO_URL
        from dvc.ui import ui
        from dvc.utils.studio import STUDIO_URL
        from dvc_studio_client.auth import StudioAuthError, get_access_token

        name = self.args.name
        hostname = self.args.hostname or os.environ.get(DVC_STUDIO_URL) or STUDIO_URL
        scopes = self.args.scopes

        studio = self.config.get("studio", {})
        if studio.get("url", hostname) == hostname and "token" in studio:
            raise DvcException(
                "Token already exists. "
                "To login with a different token, "
                "logout using 'dvc studio logout'."
            )

        open_browser = not self.args.no_open
        try:
            _, access_token = get_access_token(
                token_name=name,
                hostname=hostname,
                scopes=scopes,
                open_browser=open_browser,
                client_name="DVC",
            )
        except StudioAuthError as e:
            ui.error_write(str(e))
            return 1

        self.save_config(hostname, access_token)

        if not self.config["exp"].get("auto_push", True):
            from dvc.ui import ui

            ui.warn(
                "exp.auto_push is disabled. \n"
                "Enable with 'dvc config exp.auto_push true' "
                "to automatically push experiments to Studio."
            )

        config_path = self.config.files["global"]
        ui.write(f"Authentication complete. Saved token to {config_path}.")
        return 0

    def save_config(self, hostname, token):
        with self.config.edit("global") as conf:
            conf["studio"]["token"] = token
            conf["studio"]["url"] = hostname
            if "auto_push" not in conf["exp"]:
                conf["exp"]["auto_push"] = True


class CmdStudioLogout(CmdConfig):
    def run(self):
        from dvc.ui import ui

        with self.config.edit("global") as conf:
            if not get_in(conf, ["studio", "token"]):
                ui.error_write(
                    "Not logged in to Studio. Log in with 'dvc studio login'."
                )
                return 1

            del conf["studio"]["token"]

        ui.write("Logged out from Studio (you can log back in with 'dvc studio login')")
        return 0


class CmdStudioToken(CmdConfig):
    def run(self):
        from dvc.ui import ui

        conf = self.config.read("global")
        token = get_in(conf, ["studio", "token"])
        if not token:
            ui.error_write("Not logged in to Studio.")
            return 1

        ui.write(token)
        return 0


def add_parser(subparsers, parent_parser):
    STUDIO_HELP = "Commands to authenticate DVC with Iterative Studio"
    STUDIO_DESCRIPTION = (
        "Authenticate DVC with Studio and set the token."
        " Once this token has been properly configured,\n"
        " DVC will utilize it for seamlessly sharing live experiments\n"
        " and sending notifications to Studio regarding any experiments"
        " that have been pushed."
    )

    studio_parser = subparsers.add_parser(
        "studio",
        parents=[parent_parser],
        description=append_doc_link(STUDIO_DESCRIPTION, "studio"),
        help=STUDIO_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )
    studio_subparser = studio_parser.add_subparsers(
        dest="cmd",
        help="Use `DVC studio CMD --help` to display command-specific help.",
        required=True,
    )

    STUDIO_LOGIN_HELP = "Authenticate DVC with Studio host"
    STUDIO_LOGIN_DESCRIPTION = (
        "By default, this command authenticates the DVC with Studio\n"
        " using default scopes and assigns a random name as the token name."
    )
    login_parser = studio_subparser.add_parser(
        "login",
        parents=[parent_parser],
        description=append_doc_link(STUDIO_LOGIN_DESCRIPTION, "studio/login"),
        help=STUDIO_LOGIN_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )

    login_parser.add_argument(
        "-H",
        "--hostname",
        action="store",
        default=None,
        help="The hostname of the Studio instance to authenticate with.",
    )
    login_parser.add_argument(
        "-s",
        "--scopes",
        action="store",
        default=None,
        help="The scopes for the authentication token. ",
    )

    login_parser.add_argument(
        "-n",
        "--name",
        action="store",
        default=None,
        help="The name of the authentication token. It will be used to\n"
        "identify token shown in Studio profile.",
    )

    login_parser.add_argument(
        "--no-open",
        action="store_true",
        default=False,
        help="Use authentication flow based on user code.\n"
        "You will be presented with user code to enter in browser.\n"
        "DVC will also use this if it cannot launch browser on your behalf.",
    )
    login_parser.set_defaults(func=CmdStudioLogin)

    STUDIO_LOGOUT_HELP = "Logout user from Studio"
    STUDIO_LOGOUT_DESCRIPTION = (
        "This removes the studio token from your global config.\n"
    )

    logout_parser = studio_subparser.add_parser(
        "logout",
        parents=[parent_parser],
        description=append_doc_link(STUDIO_LOGOUT_DESCRIPTION, "studio/logout"),
        help=STUDIO_LOGOUT_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )

    logout_parser.set_defaults(func=CmdStudioLogout)

    STUDIO_TOKEN_HELP = "View the token dvc uses to contact Studio"  # noqa: S105 # nosec B105

    logout_parser = studio_subparser.add_parser(
        "token",
        parents=[parent_parser],
        description=append_doc_link(STUDIO_TOKEN_HELP, "studio/token"),
        help=STUDIO_TOKEN_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )

    logout_parser.set_defaults(func=CmdStudioToken)