fix: support 'artifact-reference' with private artifacts

Issue: #90

Author: ahal

src/taskgraph/util/parameterization.py

@@ -83,10 +83,11 @@ def repl(match):
                         f"task '{label}' has no dependency named '{dependency}'"
                     )
 
-            assert artifact_name.startswith(
-                "public/"
-            ), f"artifact-reference only supports public artifacts, not `{artifact_name}`"
-            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)
 

src/taskgraph/util/taskcluster.py

@@ -140,22 +140,9 @@ def _handle_artifact(path, response):
 
 def get_artifact_url(task_id, path, use_proxy=False):
     artifact_tmpl = liburls.api(
-        get_root_url(False), "queue", "v1", "task/{}/artifacts/{}"
+        get_root_url(use_proxy), "queue", "v1", "task/{}/artifacts/{}"
     )
-    data = artifact_tmpl.format(task_id, path)
-    if use_proxy:
-        # Until Bug 1405889 is deployed, we can't download directly
-        # from the taskcluster-proxy.  Work around by using the /bewit
-        # endpoint instead.
-        # The bewit URL is the body of a 303 redirect, which we don't
-        # want to follow (which fetches a potentially large resource).
-        response = _do_request(
-            os.environ["TASKCLUSTER_PROXY_URL"] + "/bewit",
-            data=data,
-            allow_redirects=False,
-        )
-        return response.text
-    return data
+    return artifact_tmpl.format(task_id, path)
 
 
 def get_artifact(task_id, path, use_proxy=False):

test/test_util_parameterization.py

@@ -205,6 +205,17 @@ def test_artifact_refs_in_string(assert_artifact_refs):
     )
 
 
+def test_artifact_refs_private(monkeypatch, assert_artifact_refs):
+    "resolve_task_references resolves private artifact references"
+    tc_proxy_url = "https://taskcluster-proxy.net"
+    monkeypatch.setenv("TASKCLUSTER_PROXY_URL", tc_proxy_url)
+
+    assert_artifact_refs(
+        {"artifact-reference": "<edge1/private/foo>"},
+        f"{tc_proxy_url}/api/queue/v1/task/tid1/artifacts/private/foo",
+    )
+
+
 def test_artifact_refs_self():
     "resolve_task_references raises keyerror on artifact references to `self`"
     with pytest.raises(KeyError):

test/test_util_taskcluster.py

@@ -106,6 +106,30 @@ def test_do_request(responses):
         tc._do_request("https://example.org")
 
 
+@pytest.mark.parametrize(
+    "use_proxy,expected",
+    (
+        pytest.param(
+            False,
+            "https://tc.example.com/api/queue/v1/task/abc/artifacts/public/log.txt",
+            id="use_proxy=False",
+        ),
+        pytest.param(
+            True,
+            "https://taskcluster-proxy.net/api/queue/v1/task/abc/artifacts/public/log.txt",
+            id="use_proxy=True",
+        ),
+    ),
+)
+def test_get_artifact_url(monkeypatch, use_proxy, expected):
+    if use_proxy:
+        monkeypatch.setenv("TASKCLUSTER_PROXY_URL", "https://taskcluster-proxy.net")
+
+    task_id = "abc"
+    path = "public/log.txt"
+    assert tc.get_artifact_url(task_id, path, use_proxy) == expected
+
+
 def test_get_artifact(responses, root_url):
     tid = 123
     responses.add(