diff --git a/dvc/cli/parser.py b/dvc/cli/parser.py index 88a60d1c2..31e68ad4d 100644 --- a/dvc/cli/parser.py +++ b/dvc/cli/parser.py @@ -42,6 +42,7 @@ from dvc.commands import ( repro, root, stage, + studio, unprotect, update, version, @@ -93,6 +94,7 @@ COMMANDS = [ check_ignore, data, artifacts, + studio, ] diff --git a/dvc/commands/studio.py b/dvc/commands/studio.py new file mode 100644 index 000000000..00b828557 --- /dev/null +++ b/dvc/commands/studio.py @@ -0,0 +1,175 @@ +import argparse +import os + +from funcy import get_in + +from dvc.cli.utils import append_doc_link, fix_subparsers +from dvc.commands.config import CmdConfig +from dvc.log import logger + +logger = logger.getChild(__name__) + + +class CmdStudioLogin(CmdConfig): + def run(self): + from dvc_studio_client.auth import StudioAuthError, get_access_token + + from dvc.env import DVC_STUDIO_URL + from dvc.ui import ui + from dvc.utils.studio import STUDIO_URL + + name = self.args.name + hostname = self.args.hostname or os.environ.get(DVC_STUDIO_URL) or STUDIO_URL + scopes = self.args.scopes + + try: + token_name, access_token = get_access_token( + token_name=name, + hostname=hostname, + scopes=scopes, + use_device_code=self.args.no_open, + client_name="DVC", + ) + except StudioAuthError as e: + ui.error_write(str(e)) + return 1 + + self.save_config(hostname, access_token) + ui.write( + "Authentication has been successfully completed." + "The generated token will now be accessible as" + f" {token_name} in the user's Studio profile." + ) + return 0 + + def save_config(self, hostname, token): + with self.config.edit("global") as conf: + conf["studio"]["token"] = token + conf["studio"]["url"] = hostname + + +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.") + return 1 + + del conf["studio"]["token"] + + ui.write("Logged out from Studio") + 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=argparse.RawDescriptionHelpFormatter, + ) + studio_subparser = studio_parser.add_subparsers( + dest="cmd", + help="Use `DVC studio CMD --help` to display command-specific help.", + ) + fix_subparsers(studio_subparser) + + 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=argparse.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=argparse.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=argparse.RawDescriptionHelpFormatter, + ) + + logout_parser.set_defaults(func=CmdStudioToken) diff --git a/pyproject.toml b/pyproject.toml index 0403b2d7e..6eaca3bcf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "dvc-data>=2.22.1,<2.23.0", "dvc-http>=2.29.0", "dvc-render>=0.3.1,<1", - "dvc-studio-client>=0.13.0,<1", + "dvc-studio-client>=0.17.1,<1", "dvc-task>=0.3.0,<1", "flatten_dict<1,>=0.4.1", # https://github.com/iterative/dvc/issues/9654 diff --git a/tests/unit/command/test_studio.py b/tests/unit/command/test_studio.py new file mode 100644 index 000000000..20d1b4ca0 --- /dev/null +++ b/tests/unit/command/test_studio.py @@ -0,0 +1,81 @@ +from dvc_studio_client.auth import AuthenticationExpired + +from dvc.cli import main +from dvc.utils.studio import STUDIO_URL + + +def test_studio_login_token_check_failed(mocker): + mocker.patch( + "dvc_studio_client.auth.get_access_token", side_effect=AuthenticationExpired + ) + + assert main(["studio", "login"]) == 1 + + +def test_studio_login_success(mocker, dvc): + mocker.patch( + "dvc_studio_client.auth.get_access_token", + return_value=("token_name", "isat_access_token"), + ) + + assert main(["studio", "login"]) == 0 + + config = dvc.config.load_one("global") + assert config["studio"]["token"] == "isat_access_token" + assert config["studio"]["url"] == STUDIO_URL + + +def test_studio_login_arguments(mocker): + mock = mocker.patch( + "dvc_studio_client.auth.get_access_token", + return_value=("token_name", "isat_access_token"), + ) + + assert ( + main( + [ + "studio", + "login", + "--name", + "token_name", + "--hostname", + "https://example.com", + "--scopes", + "experiments", + "--no-open", + ] + ) + == 0 + ) + + mock.assert_called_with( + token_name="token_name", + hostname="https://example.com", + scopes="experiments", + use_device_code=True, + client_name="DVC", + ) + + +def test_studio_logout(dvc): + with dvc.config.edit("global") as conf: + conf["studio"]["token"] = "isat_access_token" + + assert main(["studio", "logout"]) == 0 + config = dvc.config.load_one("global") + assert "token" not in config["studio"] + + assert main(["studio", "logout"]) == 1 + + +def test_studio_token(dvc, capsys): + with dvc.config.edit("global") as conf: + conf["studio"]["token"] = "isat_access_token" + + assert main(["studio", "token"]) == 0 + assert capsys.readouterr().out == "isat_access_token\n" + + with dvc.config.edit("global") as conf: + del conf["studio"]["token"] + + assert main(["studio", "token"]) == 1