utils.studio: require token (#9255)

This is a simple refactoring which makes the caller to choose appropriate token instead of
the studio.notify_refs making a decision for them.

dvc/repo/experiments/push.py

@@ -18,21 +18,36 @@ logger = logging.getLogger(__name__)
 
 
 def notify_refs_to_studio(scm, config, git_remote: str, **refs: List[str]):
-    token = config.get("studio_token")
     refs = compact(refs)
-    if refs and (token or config["push_exp_to_studio"]) and not env2bool("DVC_TEST"):
-        from dulwich.porcelain import get_remote_repo
+    if not refs or env2bool("DVC_TEST"):
+        return
 
-        from dvc.utils import studio
-
-        _, repo_url = get_remote_repo(scm.dulwich.repo, git_remote)
-        studio_url = config.get("studio_url")
-        studio.notify_refs(
-            repo_url,
-            default_token=token,
-            studio_url=studio_url,
-            **refs,  # type: ignore[arg-type]
+    if not (config.get("studio_token") or config["push_exp_to_studio"]):
+        logger.debug(
+            "Either feature.studio_token or feature.push_exp_to_studio config "
+            "needs to be set."
         )
+        return
+
+    import os
+
+    token = os.environ.get("STUDIO_TOKEN") or config.get("studio_token")
+    if not token:
+        logger.debug("Studio token not found.")
+        return
+
+    from dulwich.porcelain import get_remote_repo
+
+    from dvc.utils import studio
+
+    _, repo_url = get_remote_repo(scm.dulwich.repo, git_remote)
+    studio_url = config.get("studio_url")
+    studio.notify_refs(
+        repo_url,
+        token,
+        studio_url=studio_url,
+        **refs,  # type: ignore[arg-type]
+    )
 
 
 @locked

dvc/utils/studio.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from urllib.parse import urljoin
 
@@ -13,7 +12,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 STUDIO_URL = "https://studio.iterative.ai"
-STUDIO_TOKEN = "STUDIO_TOKEN"  # noqa: S105
 
 
 def post(
@@ -38,9 +36,9 @@ def post(
 
 def notify_refs(
     repo_url: str,
+    token: str,
     *,
-    default_token: Optional[str] = None,
-    studio_url: Optional[str] = None,
+    studio_url: Optional[str] = STUDIO_URL,
     **refs: List[str],
 ) -> None:
     extra_keys = refs.keys() - {"pushed", "removed"}
@@ -50,11 +48,6 @@ def notify_refs(
     if not refs:
         return
 
-    token = os.getenv(STUDIO_TOKEN) or default_token
-    if not token:
-        logger.debug("Studio token not found.")
-        return
-
     logger.debug(
         "notifying Studio%s about updated experiments",
         f" ({studio_url})" if studio_url else "",
@@ -62,7 +55,7 @@ def notify_refs(
     data = {"repo_url": repo_url, "client": "dvc", "refs": refs}
 
     try:
-        post("/webhook/dvc", token=token, data=data, url=studio_url)
+        post("/webhook/dvc", token, data, url=studio_url)
     except requests.RequestException as e:
         logger.debug("", exc_info=True)
 

tests/unit/utils/test_studio.py

@@ -22,7 +22,7 @@ def test_notify_refs(mocker, status_code):
 
     notify_refs(
         "[email protected]:iterative/dvc.git",
-        default_token="TOKEN",
+        "TOKEN",
         pushed=["p1", "p2"],
         removed=["r1", "r2"],
     )