From 74e1af942c972c63bd7301d0796b5c3bf4c036ff Mon Sep 17 00:00:00 2001
From: abhishekmadan30 <amadan@mozilla.com>
Date: Tue, 7 Oct 2025 17:43:15 -0400
Subject: [PATCH] fix: added use_proxy flag to get_artifact_url

---
 src/taskgraph/util/parameterization.py |  6 ++++-
 src/taskgraph/util/taskcluster.py      | 22 ++++++++++++++----
 test/test_util_taskcluster.py          | 32 +++++++++++++++++++++-----
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/taskgraph/util/parameterization.py b/src/taskgraph/util/parameterization.py
index 0b814cffb..39cbb2ea2 100644
--- a/src/taskgraph/util/parameterization.py
+++ b/src/taskgraph/util/parameterization.py
@@ -90,7 +90,11 @@ def repl(match):
                         f"task '{label}' has no dependency named '{dependency}'"
                     )
 
-            return get_artifact_url(task_id, artifact_name)
+            use_proxy = False
+            if not artifact_name.startswith("public/"):
+                use_proxy = True
+
+            return get_artifact_url(task_id, artifact_name, use_proxy=use_proxy)
 
         return ARTIFACT_REFERENCE_PATTERN.sub(repl, val)
 
diff --git a/src/taskgraph/util/taskcluster.py b/src/taskgraph/util/taskcluster.py
index 9878a464f..49fff7234 100644
--- a/src/taskgraph/util/taskcluster.py
+++ b/src/taskgraph/util/taskcluster.py
@@ -33,8 +33,8 @@
 
 
 @functools.lru_cache(maxsize=None)
-def get_root_url():
-    if "TASKCLUSTER_PROXY_URL" in os.environ:
+def get_root_url(block_proxy=False):
+    if "TASKCLUSTER_PROXY_URL" in os.environ and not block_proxy:
         logger.debug(
             "Using taskcluster-proxy at {}".format(os.environ["TASKCLUSTER_PROXY_URL"])
         )
@@ -140,8 +140,22 @@ def get_session():
     return requests_retry_session(retries=5)
 
 
-def get_artifact_url(task_id, path):
-    artifact_tmpl = liburls.api(get_root_url(), "queue", "v1", "task/{}/artifacts/{}")
+def get_artifact_url(task_id, path, use_proxy=False):
+    if use_proxy:
+        try:
+            url = liburls.normalize_root_url(os.environ["TASKCLUSTER_PROXY_URL"])
+        except KeyError:
+            if "TASK_ID" not in os.environ:
+                raise RuntimeError(
+                    "taskcluster-proxy is not available when not executing in a task"
+                )
+            else:
+                raise RuntimeError("taskcluster-proxy is not enabled for this task")
+
+    else:
+        url = get_root_url(block_proxy=True)
+
+    artifact_tmpl = liburls.api(url, "queue", "v1", "task/{}/artifacts/{}")
     return artifact_tmpl.format(task_id, path)
 
 
diff --git a/test/test_util_taskcluster.py b/test/test_util_taskcluster.py
index 59b840c68..7bbcd3584 100644
--- a/test/test_util_taskcluster.py
+++ b/test/test_util_taskcluster.py
@@ -67,21 +67,41 @@ def test_get_root_url(monkeypatch):
 
 
 def test_get_artifact_url(monkeypatch):
+    tc.get_root_url.cache_clear()
     task_id = "abc"
     path = "public/log.txt"
+    expected = "https://tc.example.com/api/queue/v1/task/abc/artifacts/public/log.txt"
+    expected_proxy = (
+        "https://taskcluster-proxy.net/api/queue/v1/task/abc/artifacts/public/log.txt"
+    )
 
-    # Test with default root URL
+    # Test with default root URL (no proxy)
     tc.get_root_url.cache_clear()
-    expected = "https://tc.example.com/api/queue/v1/task/abc/artifacts/public/log.txt"
     assert tc.get_artifact_url(task_id, path) == expected
 
-    # Test with proxy URL (takes priority)
+    # Test that proxy URL is NOT used by default (since use_proxy defaults to False)
     monkeypatch.setenv("TASKCLUSTER_PROXY_URL", "https://taskcluster-proxy.net")
     tc.get_root_url.cache_clear()
-    expected_proxy = (
-        "https://taskcluster-proxy.net/api/queue/v1/task/abc/artifacts/public/log.txt"
+    assert tc.get_artifact_url(task_id, path) == expected
+
+    # Test with use_proxy=True (should use proxy URL)
+    assert tc.get_artifact_url(task_id, path, use_proxy=True) == expected_proxy
+
+    # Test with use_proxy=True but no proxy available (not in a task)
+    monkeypatch.delenv("TASKCLUSTER_PROXY_URL")
+    monkeypatch.delenv("TASK_ID", raising=False)
+    tc.get_root_url.cache_clear()
+    with pytest.raises(RuntimeError) as exc:
+        tc.get_artifact_url(task_id, path, use_proxy=True)
+    assert "taskcluster-proxy is not available when not executing in a task" in str(
+        exc.value
     )
-    assert tc.get_artifact_url(task_id, path) == expected_proxy
+
+    # Test with use_proxy=True but proxy not enabled (in a task without proxy)
+    monkeypatch.setenv("TASK_ID", "some-task-id")
+    with pytest.raises(RuntimeError) as exc:
+        tc.get_artifact_url(task_id, path, use_proxy=True)
+    assert "taskcluster-proxy is not enabled for this task" in str(exc.value)
 
 
 def test_get_artifact(responses, root_url):