studio: Add Studio authentication to DVC

dvc/cli/parser.py

@@ -42,6 +42,7 @@
     repro,
     root,
     stage,
+    studio,
     unprotect,
     update,
     version,
@@ -93,6 +94,7 @@
     check_ignore,
     data,
     artifacts,
+    studio,
 ]
 
 

dvc/commands/studio.py

@@ -0,0 +1,176 @@
+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.use_device_code,
+                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(
+        "-d",
+        "--use-device-code",
+        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)

dvc/utils/studio.py

@@ -35,7 +35,7 @@ def post(
 
     logger.trace("Sending %s to %s", data, url)
 
-    headers = {"Authorization": f"token {token}"}
+    headers = {"Authorization": f"token {token}"} if token else None
     r = session.post(
         url, json=data, headers=headers, timeout=timeout, allow_redirects=False
     )

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

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",
+                "--use-device-code",
+            ]
+        )
+        == 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