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,172 @@
+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__)
+
+DEFAULT_SCOPES = "live,dvc_experiment,view_url,dql,download_model"
+
+
+class CmdStudioLogin(CmdConfig):
+    def run(self):
+        from dvc_studio_client.auth import StudioAuthError, initiate_authorization
+
+        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 or DEFAULT_SCOPES
+
+        try:
+            token_name, access_token = initiate_authorization(
+                name=name,
+                hostname=hostname,
+                scopes=scopes,
+                use_device_code=self.args.use_device_code,
+            )
+        except StudioAuthError as e:
+            ui.error_write(str(e))
+            return 1
+
+        self.save_config(hostname, access_token)
+        ui.write(
+            f"Authentication successful. The token will be "
+            f"available as {token_name} in 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 = "Authenticate dvc with Iterative Studio"
+    STUDIO_DESCRIPTION = (
+        "Authorize dvc with Studio and set the token. When this is\n"
+        "set, DVC uses this to share live experiments and notify\n"
+        "Studio about pushed experiments."
+    )
+
+    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 authorize dvc with Studio with\n"
+        " default scopes and a random  name as 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 command helps to log out user from DVC Studio.\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
     )
@@ -66,7 +66,7 @@ def notify_refs(
     try:
         r = post("webhook/dvc", token, data, base_url=base_url)
     except requests.RequestException as e:
-        logger.trace("", exc_info=True)
+        logger.trace("", exc_info=True)  # type: ignore[attr-defined]
 
         msg = str(e)
         if e.response is None:

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "dvc-data>=2.22.0,<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.16.1,<1",
     "dvc-task>=0.3.0,<1",
     "flatten_dict<1,>=0.4.1",
     # https://github.com/iterative/dvc/issues/9654

tests/func/test_studio.py

@@ -0,0 +1,105 @@
+from requests import Response
+
+from dvc.cli import main
+from dvc.commands.studio import DEFAULT_SCOPES
+from dvc.utils.studio import STUDIO_URL
+from tests.unit.command.test_studio import MOCK_RESPONSE
+
+
+def test_auth_expired(mocker, M):
+    mock_login_post = mocker.patch(
+        "requests.post", return_value=_mock_response(mocker, 200, MOCK_RESPONSE)
+    )
+    mock_poll_post = mocker.patch(
+        "requests.Session.post",
+        side_effect=[
+            _mock_response(mocker, 400, {"detail": "authorization_expired"}),
+        ],
+    )
+
+    assert main(["studio", "login"]) == 1
+
+    assert mock_login_post.call_args == mocker.call(
+        url=f"{STUDIO_URL}/api/device-login",
+        json={
+            "client_name": "dvc",
+            "scopes": DEFAULT_SCOPES.split(","),
+        },
+        headers={"Content-type": "application/json"},
+        timeout=5,
+    )
+
+    assert mock_poll_post.call_args_list == [
+        mocker.call(
+            f"{STUDIO_URL}/api/device-login/token",
+            json={"code": "random-value"},
+            timeout=5,
+            allow_redirects=False,
+        ),
+    ]
+
+
+def test_studio_success(mocker, dvc):
+    mocker.patch("time.sleep")
+    mock_login_post = mocker.patch(
+        "requests.post", return_value=_mock_response(mocker, 200, MOCK_RESPONSE)
+    )
+    mock_poll_post = mocker.patch(
+        "requests.Session.post",
+        side_effect=[
+            _mock_response(mocker, 400, {"detail": "authorization_pending"}),
+            _mock_response(mocker, 200, {"access_token": "isat_access_token"}),
+        ],
+    )
+
+    assert (
+        main(
+            [
+                "studio",
+                "login",
+                "--name",
+                "token_name",
+                "--hostname",
+                "https://example.com",
+                "--scopes",
+                "live",
+            ]
+        )
+        == 0
+    )
+
+    assert mock_login_post.call_args_list == [
+        mocker.call(
+            url="https://example.com/api/device-login",
+            json={"client_name": "dvc", "token_name": "token_name", "scopes": ["live"]},
+            headers={"Content-type": "application/json"},
+            timeout=5,
+        )
+    ]
+    assert mock_poll_post.call_count == 2
+    assert mock_poll_post.call_args_list == [
+        mocker.call(
+            f"{STUDIO_URL}/api/device-login/token",
+            json={"code": "random-value"},
+            timeout=5,
+            allow_redirects=False,
+        ),
+        mocker.call(
+            f"{STUDIO_URL}/api/device-login/token",
+            json={"code": "random-value"},
+            timeout=5,
+            allow_redirects=False,
+        ),
+    ]
+
+    config = dvc.config.load_one("global")
+    assert config["studio"]["token"] == "isat_access_token"
+    assert config["studio"]["url"] == "https://example.com"
+
+
+def _mock_response(mocker, status_code, json):
+    response = Response()
+    response.status_code = status_code
+    mocker.patch.object(response, "json", side_effect=[json])
+
+    return response