Repository: iterative/dvc
Pull request #9254 (2023-03-28 06:32:41)

use git remote to find out remote repo url when pinging Studio


Patch:
diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py
index 12a7dd32c..677a1c468 100644
--- a/dvc/repo/experiments/push.py
+++ b/dvc/repo/experiments/push.py
@@ -17,15 +17,18 @@ from .utils import exp_commits, exp_refs, exp_refs_by_baseline, resolve_name
 logger = logging.getLogger(__name__)
 
 
-def notify_refs_to_studio(config, git_remote: str, **refs: List[str]):
+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
+
         from dvc.utils import studio
 
+        _, repo_url = get_remote_repo(scm.dulwich.repo, git_remote)
         studio_url = config.get("studio_url")
         studio.notify_refs(
-            git_remote,
+            repo_url,
             default_token=token,
             studio_url=studio_url,
             **refs,  # type: ignore[arg-type]
@@ -88,7 +91,9 @@ def push(  # noqa: C901
 
     refs = push_result[SyncStatus.SUCCESS]
     pushed_refs = [str(r) for r in refs]
-    notify_refs_to_studio(repo.config["feature"], git_remote, pushed=pushed_refs)
+    notify_refs_to_studio(
+        repo.scm, repo.config["feature"], git_remote, pushed=pushed_refs
+    )
     return [ref.name for ref in refs]
 
 
diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py
index 27675dc2d..74cb64ae0 100644
--- a/dvc/repo/experiments/remove.py
+++ b/dvc/repo/experiments/remove.py
@@ -79,7 +79,9 @@ def remove(  # noqa: C901, PLR0912
         from .push import notify_refs_to_studio
 
         removed_refs = [str(r) for r in exp_ref_list]
-        notify_refs_to_studio(repo.config["feature"], git_remote, removed=removed_refs)
+        notify_refs_to_studio(
+            repo.scm, repo.config["feature"], git_remote, removed=removed_refs
+        )
     return removed
 
 
diff --git a/dvc/utils/studio.py b/dvc/utils/studio.py
index 071f1bb0d..ecc1a51b8 100644
--- a/dvc/utils/studio.py
+++ b/dvc/utils/studio.py
@@ -1,10 +1,9 @@
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from urllib.parse import urljoin
 
 import requests
-from dvc_studio_client.post_live_metrics import get_studio_repo_url
 from funcy import compact
 from requests.adapters import HTTPAdapter
 
@@ -14,20 +13,9 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 STUDIO_URL = "https://studio.iterative.ai"
-STUDIO_REPO_URL = "STUDIO_REPO_URL"
 STUDIO_TOKEN = "STUDIO_TOKEN"  # noqa: S105
 
 
-def get_studio_token_and_repo_url(
-    default_token: Optional[str] = None,
-    repo_url_finder: Optional[Callable[[], Optional[str]]] = None,
-) -> Tuple[Optional[str], Optional[str]]:
-    token = os.getenv(STUDIO_TOKEN) or default_token
-    url_finder = repo_url_finder or get_studio_repo_url
-    repo_url = os.getenv(STUDIO_REPO_URL) or url_finder()
-    return token, repo_url
-
-
 def post(
     endpoint: str,
     token: str,
@@ -49,36 +37,24 @@ def post(
 
 
 def notify_refs(
-    git_remote: str,
+    repo_url: str,
+    *,
     default_token: Optional[str] = None,
     studio_url: Optional[str] = None,
-    repo_url_finder: Optional[Callable[[], Optional[str]]] = None,
     **refs: List[str],
 ) -> None:
-    # TODO: Should we use git_remote to associate with Studio project
-    # instead of using `git ls-remote` on fallback?
+    extra_keys = refs.keys() - {"pushed", "removed"}
+    assert not extra_keys, f"got extra args: {extra_keys}"
+
     refs = compact(refs)
     if not refs:
         return
 
-    assert git_remote
-    token, repo_url = get_studio_token_and_repo_url(
-        default_token=default_token,
-        repo_url_finder=repo_url_finder,
-    )
+    token = os.getenv(STUDIO_TOKEN) or default_token
     if not token:
         logger.debug("Studio token not found.")
         return
 
-    if not repo_url:
-        logger.warning(
-            "Could not detect repository url. "
-            "Please set %s environment variable "
-            "to your remote git repository url. ",
-            STUDIO_REPO_URL,
-        )
-        return
-
     logger.debug(
         "notifying Studio%s about updated experiments",
         f" ({studio_url})" if studio_url else "",


Test patch:
diff --git a/tests/unit/utils/test_studio.py b/tests/unit/utils/test_studio.py
index df9598764..c3767c898 100644
--- a/tests/unit/utils/test_studio.py
+++ b/tests/unit/utils/test_studio.py
@@ -21,9 +21,8 @@ def test_notify_refs(mocker, status_code):
     mock_post = mocker.patch("requests.Session.post", return_value=response)
 
     notify_refs(
-        "origin",
-        "TOKEN",
-        repo_url_finder=lambda: "git@github.com:iterative/dvc.git",
+        "git@github.com:iterative/dvc.git",
+        default_token="TOKEN",
         pushed=["p1", "p2"],
         removed=["r1", "r2"],
     )