Skip to content

Commit

Permalink
utils.studio: disallow redirects, misc API refactoring (#9301)
Browse files Browse the repository at this point in the history
  • Loading branch information
skshetry authored Apr 4, 2023
1 parent 071a5b7 commit 8be71be
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 7 deletions.
16 changes: 9 additions & 7 deletions dvc/utils/studio.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,21 +15,23 @@


def post(
endpoint: str,
url: str,
token: str,
data: Dict[str, Any],
url: Optional[str] = STUDIO_URL,
base_url: Optional[str] = STUDIO_URL,
max_retries: int = 3,
timeout: int = 5,
) -> "Response":
endpoint = urljoin(url or STUDIO_URL, endpoint)
url = urljoin(base_url or STUDIO_URL, url)
session = requests.Session()
session.mount(endpoint, HTTPAdapter(max_retries=max_retries))
session.mount(url, HTTPAdapter(max_retries=max_retries))

logger.trace("Sending %s to %s", data, endpoint) # type: ignore[attr-defined]
logger.trace("Sending %s to %s", data, url) # type: ignore[attr-defined]

headers = {"Authorization": f"token {token}"}
r = session.post(endpoint, json=data, headers=headers, timeout=timeout)
r = session.post(
url, json=data, headers=headers, timeout=timeout, allow_redirects=False
)
r.raise_for_status()
return r

Expand All @@ -55,7 +57,7 @@ def notify_refs(
data = {"repo_url": repo_url, "client": "dvc", "refs": refs}

try:
r = post("/webhook/dvc", token, data, url=studio_url)
r = post("/webhook/dvc", token, data, base_url=studio_url)
except requests.RequestException as e:
logger.trace("", exc_info=True) # type: ignore[attr-defined]

Expand Down
1 change: 1 addition & 0 deletions tests/unit/utils/test_studio.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,4 +41,5 @@ def test_notify_refs(mocker, status_code, side_effect):
},
headers={"Authorization": "token TOKEN"},
timeout=5,
allow_redirects=False,
)

0 comments on commit 8be71be

Please sign in to comment.