From ef507efa04479b0595fcf502a04a8a883d272dc5 Mon Sep 17 00:00:00 2001
From: Dave Berenbaum <dave@iterative.ai>
Date: Fri, 10 May 2024 09:59:20 -0400
Subject: [PATCH] artifacts: get studio info from local config

---
 dvc/repo/artifacts.py                  | 24 ++++++++++++++++--------
 tests/func/artifacts/test_artifacts.py |  9 +++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/dvc/repo/artifacts.py b/dvc/repo/artifacts.py
index dda6c4ef3..43cd80732 100644
--- a/dvc/repo/artifacts.py
+++ b/dvc/repo/artifacts.py
@@ -17,7 +17,7 @@ from dvc.utils.objects import cached_property
 from dvc.utils.serialize import modify_yaml
 
 if TYPE_CHECKING:
-    from gto.tag import Tag as GTOTag
+    from gto.tag import Tag as GTOTag  # type: ignore[import-not-found]
     from scmrepo.git import GitTag
 
     from dvc.repo import Repo
@@ -27,8 +27,8 @@ logger = logger.getChild(__name__)
 
 
 def check_name_format(name: str) -> None:
-    from gto.constants import assert_name_is_valid
-    from gto.exceptions import ValidationError
+    from gto.constants import assert_name_is_valid  # type: ignore[import-not-found]
+    from gto.exceptions import ValidationError  # type: ignore[import-not-found]
 
     try:
         assert_name_is_valid(name)
@@ -132,7 +132,7 @@ class Artifacts:
         self, name: str, version: Optional[str] = None, stage: Optional[str] = None
     ):
         """Return revision containing the given artifact."""
-        from gto.base import sort_versions
+        from gto.base import sort_versions  # type: ignore[import-not-found]
         from gto.tag import find, parse_tag
 
         assert not (version and stage)
@@ -314,6 +314,13 @@ class Artifacts:
 
         name = _reformat_name(name)
         saved_exc: Optional[Exception] = None
+
+        local_dvc_studio_config = Config().get("studio", {})
+        args_dvc_studio_config = {}
+        if config and not isinstance(config, dict):
+            config = Config.load_file(config)
+            args_dvc_studio_config = config.get("studio", {})
+
         try:
             logger.trace("Trying studio-only config")
             return cls._download_studio(
@@ -324,14 +331,13 @@ class Artifacts:
                 out=out,
                 force=force,
                 jobs=jobs,
+                dvc_studio_config=local_dvc_studio_config | args_dvc_studio_config,
             )
         except FileExistsLocallyError:
             raise
         except Exception as exc:  # noqa: BLE001
             saved_exc = exc
 
-        if config and not isinstance(config, dict):
-            config = Config.load_file(config)
         with Repo.open(
             url=url,
             subrepos=True,
@@ -341,7 +347,7 @@ class Artifacts:
             remote_config=remote_config,
         ) as repo:
             logger.trace("Trying repo [studio] config")
-            dvc_studio_config = dict(repo.config.get("studio"))
+            repo_dvc_studio_config = repo.config.get("studio", {})
             try:
                 return cls._download_studio(
                     url,
@@ -351,7 +357,9 @@ class Artifacts:
                     out=out,
                     force=force,
                     jobs=jobs,
-                    dvc_studio_config=dvc_studio_config,
+                    dvc_studio_config=local_dvc_studio_config
+                    | repo_dvc_studio_config
+                    | args_dvc_studio_config,
                 )
             except FileExistsLocallyError:
                 raise
diff --git a/tests/func/artifacts/test_artifacts.py b/tests/func/artifacts/test_artifacts.py
index 12456b317..5b539c22d 100644
--- a/tests/func/artifacts/test_artifacts.py
+++ b/tests/func/artifacts/test_artifacts.py
@@ -281,3 +281,12 @@ def test_artifacts_download_subrepo(tmp_dir, scm, sub):
     assert Artifacts.get(subdir.fs_path, name, force=True) == result
     with subdir.chdir():
         assert Artifacts.get(".", name, force=True) == result
+
+
+def test_artifacts_download_studio(tmp_dir, dvc, mocker):
+    with dvc.config.edit("global") as conf:
+        conf["studio"]["token"] = "mytoken"
+
+    download_studio = mocker.patch("dvc.repo.artifacts.Artifacts._download_studio")
+    Artifacts.get("myart.pkl", "myart.pkl")
+    assert download_studio.call_args.kwargs["dvc_studio_config"]["token"] == "mytoken"