fix: replaced taskcluster rest api calls with package

fixes #221

Author: abhishekmadan30

pyproject.toml

@@ -29,6 +29,7 @@ dependencies = [
   "redo>=2.0",
   "requests>=2.25",
   "slugid>=2.0",
+  "taskcluster>=55.0",
   "taskcluster-urls>=11.0",
   "voluptuous>=0.12.1",
 ]

src/taskgraph/actions/retrigger.py

@@ -163,7 +163,8 @@ def retrigger_action(parameters, graph_config, input, task_group_id, task_id):
         to_run = full_task_graph.graph.transitive_closure(
             set(to_run), reverse=True
         ).nodes
-        to_run = to_run & set(label_to_taskid.keys())
+        if label_to_taskid:
+            to_run = to_run & set(label_to_taskid.keys())
         with_downstream = " (with downstream) "
 
     times = input.get("times", 1)
@@ -202,7 +203,7 @@ def rerun_action(parameters, graph_config, input, task_group_id, task_id):
         parameters, graph_config, task_group_id=task_group_id
     )
     label = task["metadata"]["name"]
-    if task_id not in label_to_taskid.values():
+    if label_to_taskid and task_id not in label_to_taskid.values():
         logger.error(
             f"Refusing to rerun {label}: taskId {task_id} not in decision task {decision_task_id} label_to_taskid!"
         )
@@ -276,7 +277,8 @@ def retrigger_multiple(parameters, graph_config, input, task_group_id, task_id):
             # In practice, this shouldn't matter, as only completed tasks
             # are pulled in from other pushes and treeherder won't pass
             # those labels.
-            _rerun_task(label_to_taskid[label], label)
+            if label_to_taskid and label in label_to_taskid:
+                _rerun_task(label_to_taskid[label], label)
 
         for j in range(times):
             suffix = f"{i}-{j}"

src/taskgraph/create.py

@@ -4,15 +4,14 @@
 
 
 import logging
-import os
 import sys
 from concurrent import futures
 
 from slugid import nice as slugid
 
 from taskgraph.util import json
 from taskgraph.util.parameterization import resolve_timestamps
-from taskgraph.util.taskcluster import CONCURRENCY, get_session
+from taskgraph.util.taskcluster import CONCURRENCY, get_session, get_taskcluster_queue
 from taskgraph.util.time import current_json_time
 
 logger = logging.getLogger(__name__)
@@ -104,9 +103,6 @@ def submit(task_id, label, task_def):
 
 
 def create_task(session, task_id, label, task_def):
-    # create the task using 'http://taskcluster/queue', which is proxied to the queue service
-    # with credentials appropriate to this task.
-
     # Resolve timestamps
     now = current_json_time(datetime_format=True)
     task_def = resolve_timestamps(now, task_def)
@@ -123,16 +119,5 @@ def create_task(session, task_id, label, task_def):
         return
 
     logger.info(f"Creating task with taskId {task_id} for {label}")
-    proxy_url = os.environ.get("TASKCLUSTER_PROXY_URL", "http://taskcluster").rstrip(
-        "/"
-    )
-    res = session.put(
-        f"{proxy_url}/queue/v1/task/{task_id}",
-        json=task_def,
-    )
-    if res.status_code != 200:
-        try:
-            logger.error(res.json()["message"])
-        except Exception:
-            logger.error(res.text)
-        res.raise_for_status()
+    queue = get_taskcluster_queue()
+    queue.createTask(task_id, task_def)

src/taskgraph/run-task/fetch-content

@@ -39,6 +39,11 @@ try:
 except ImportError:
     certifi = None
 
+try:
+    import taskcluster
+except ImportError:
+    taskcluster = None
+
 
 CONCURRENCY = multiprocessing.cpu_count()
 
@@ -734,15 +739,30 @@ def git_checkout_archive(
         env = os.environ.copy()
         keypath = ""
         if ssh_key:
-            taskcluster_secret_url = api(
-                os.environ.get("TASKCLUSTER_PROXY_URL"),
-                "secrets",
-                "v1",
-                "secret/{keypath}".format(keypath=ssh_key),
-            )
-            taskcluster_secret = b"".join(stream_download(taskcluster_secret_url))
-            taskcluster_secret = json.loads(taskcluster_secret)
-            sshkey = taskcluster_secret["secret"]["ssh_privkey"]
+            err = False
+
+            if taskcluster:
+                secrets_client = get_taskcluster_secrets()
+                if secrets_client:
+                    taskcluster_secret = secrets_client.get(ssh_key)
+                    if taskcluster_secret:
+                        sshkey = taskcluster_secret["secret"]["ssh_privkey"]
+                    else:
+                        err=True
+                else:
+                    err = True
+
+            if not taskcluster or err:
+                # Fallback to old method if taskcluster package not available
+                taskcluster_secret_url = api(
+                    os.environ.get("TASKCLUSTER_PROXY_URL"),
+                    "secrets",
+                    "v1",
+                    "secret/{keypath}".format(keypath=ssh_key),
+                )
+                taskcluster_secret = b"".join(stream_download(taskcluster_secret_url))
+                taskcluster_secret = json.loads(taskcluster_secret)
+                sshkey = taskcluster_secret["secret"]["ssh_privkey"]
 
             keypath = temp_dir.joinpath("ssh-key")
             keypath.write_text(sshkey)
@@ -900,6 +920,19 @@ def api(root_url, service, version, path):
     )
 
 
+def get_taskcluster_secrets():
+    """Get a taskcluster secrets client."""
+    if not taskcluster:
+        return None
+    
+    if "TASKCLUSTER_PROXY_URL" in os.environ:
+        secrets_options = {"rootUrl": os.environ["TASKCLUSTER_PROXY_URL"]}
+    else:
+        secrets_options = taskcluster.optionsFromEnvironment()
+
+    return taskcluster.Secrets(secrets_options)
+
+
 def get_hash(fetch, root_url):
     path = "task/{task}/artifacts/{artifact}".format(
         task=fetch["task"], artifact="public/chain-of-trust.json"