Skip to content

Commit

Permalink
utils.studio: require token (#9255)
Browse files Browse the repository at this point in the history
This is a simple refactoring which makes the caller to choose appropriate token instead of
the studio.notify_refs making a decision for them.
  • Loading branch information
skshetry authored and daavoo committed Mar 28, 2023
1 parent b03f8d6 commit 36af16e
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 24 deletions.
41 changes: 28 additions & 13 deletions dvc/repo/experiments/push.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,21 +18,36 @@


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(
repo_url,
default_token=token,
studio_url=studio_url,
**refs, # type: ignore[arg-type]
if not refs or env2bool("DVC_TEST"):
return

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
Expand Down
13 changes: 3 additions & 10 deletions dvc/utils/studio.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import logging
import os
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from urllib.parse import urljoin

Expand All @@ -13,7 +12,6 @@
logger = logging.getLogger(__name__)

STUDIO_URL = "https://studio.iterative.ai"
STUDIO_TOKEN = "STUDIO_TOKEN" # noqa: S105


def post(
Expand All @@ -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"}
Expand All @@ -50,19 +48,14 @@ 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 "",
)
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)

Expand Down
2 changes: 1 addition & 1 deletion tests/unit/utils/test_studio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
)
Expand Down

0 comments on commit 36af16e

Please sign in to comment.