fix: replaced taskcluster rest api calls with package

fix: replaced taskcluster rest api calls with package

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/cancel.py

@@ -26,7 +26,7 @@ def cancel_action(parameters, graph_config, input, task_group_id, task_id):
     # Note that this is limited by the scopes afforded to generic actions to
     # only cancel tasks with the level-specific schedulerId.
     try:
-        cancel_task(task_id, use_proxy=True)
+        cancel_task(task_id)
     except requests.HTTPError as e:
         if e.response.status_code == 409:
             # A 409 response indicates that this task is past its deadline.  It

src/taskgraph/actions/cancel_all.py

@@ -35,7 +35,7 @@ def cancel_all_action(parameters, graph_config, input, task_group_id, task_id):
     def do_cancel_task(task_id):
         logger.info(f"Cancelling task {task_id}")
         try:
-            cancel_task(task_id, use_proxy=True)
+            cancel_task(task_id)
         except requests.HTTPError as e:
             if e.response.status_code == 409:
                 # A 409 response indicates that this task is past its deadline.  It

src/taskgraph/actions/retrigger.py

@@ -147,7 +147,7 @@ def retrigger_action(parameters, graph_config, input, task_group_id, task_id):
     )
 
     task = taskcluster.get_task_definition(task_id)
-    label = task["metadata"]["name"]
+    label = task["metadata"]["name"]  # type: ignore
 
     with_downstream = " "
     to_run = [label]
@@ -201,7 +201,7 @@ def rerun_action(parameters, graph_config, input, task_group_id, task_id):
     decision_task_id, full_task_graph, label_to_taskid = fetch_graph_and_labels(
         parameters, graph_config, task_group_id=task_group_id
     )
-    label = task["metadata"]["name"]
+    label = task["metadata"]["name"]  # type: ignore
     if 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!"

src/taskgraph/actions/util.py

@@ -56,12 +56,10 @@ def fetch_graph_and_labels(parameters, graph_config, task_group_id=None):
         # for old ones
         def fetch_action(task_id):
             logger.info(f"fetching label-to-taskid.json for action task {task_id}")
-            try:
-                run_label_to_id = get_artifact(task_id, "public/label-to-taskid.json")
+            run_label_to_id = get_artifact(task_id, "public/label-to-taskid.json")
+            if label_to_taskid and run_label_to_id:
                 label_to_taskid.update(run_label_to_id)
-            except HTTPError as e:
-                if e.response.status_code != 404:
-                    raise
+            else:
                 logger.debug(f"No label-to-taskid.json found for {task_id}: {e}")
 
         # for backwards compatibility, look up actions via pushlog-id
@@ -84,7 +82,7 @@ def fetch_cron(task_id):
             logger.info(f"fetching label-to-taskid.json for cron task {task_id}")
             try:
                 run_label_to_id = get_artifact(task_id, "public/label-to-taskid.json")
-                label_to_taskid.update(run_label_to_id)
+                label_to_taskid.update(run_label_to_id)  # type: ignore
             except HTTPError as e:
                 if e.response.status_code != 404:
                     raise

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_client
 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_client("queue")
+    queue.createTask(task_id, task_def)

src/taskgraph/docker.py

@@ -425,9 +425,10 @@ def load_task(
 
     if "payload" not in task_def or not (image := task_def["payload"].get("image")):
         print("Tasks without a `payload.image` are not supported!")
+
         return 1
 
-    command = task_def["payload"].get("command")
+    command = task_def["payload"].get("command")  # type: ignore
     if not command or not command[0].endswith("run-task"):
         print("Only tasks using `run-task` are supported!")
         return 1
@@ -470,12 +471,12 @@ def load_task(
     # Set some env vars the worker would normally set.
     env = {
         "RUN_ID": "0",
-        "TASK_GROUP_ID": task_def.get("taskGroupId", ""),
+        "TASK_GROUP_ID": task_def.get("taskGroupId", ""),  # type: ignore
         "TASK_ID": task_id,
-        "TASKCLUSTER_ROOT_URL": get_root_url(False),
+        "TASKCLUSTER_ROOT_URL": get_root_url(),
     }
     # Add the task's environment variables.
-    env.update(task_def["payload"].get("env", {}))
+    env.update(task_def["payload"].get("env", {}))  # type: ignore
 
     # run-task expects the worker to mount a volume for each path defined in
     # TASKCLUSTER_CACHES, delete them to avoid needing to do the same.

src/taskgraph/optimize/strategies.py

@@ -48,14 +48,15 @@ def should_replace_task(self, task, params, deadline, arg):
                     status = status_task(task_id)
                 # status can be `None` if we're in `testing` mode
                 # (e.g. test-action-callback)
-                if not status or status.get("state") in ("exception", "failed"):
+                if not status or status.get("state") in ("exception", "failed"):  # type: ignore
                     logger.debug(
                         f"not replacing {task.label} with {task_id} because it is in failed or exception state"
                     )
                     continue
 
                 if deadline and datetime.strptime(
-                    status["expires"], self.fmt
+                    status["expires"],  # type: ignore
+                    self.fmt,
                 ) < datetime.strptime(deadline, self.fmt):
                     logger.debug(
                         f"not replacing {task.label} with {task_id} because it expires before {deadline}"

src/taskgraph/transforms/run/run_task.py

@@ -183,9 +183,8 @@ def script_url(config, script):
     if "MOZ_AUTOMATION" in os.environ and "TASK_ID" not in os.environ:
         raise Exception("TASK_ID must be defined to use run-task on generic-worker")
     task_id = os.environ.get("TASK_ID", "<TASK_ID>")
-    # use_proxy = False to avoid having all generic-workers turn on proxy
     # Assumes the cluster allows anonymous downloads of public artifacts
-    tc_url = taskcluster.get_root_url(False)
+    tc_url = taskcluster.get_root_url()
     # TODO: Use util/taskcluster.py:get_artifact_url once hack for Bug 1405889 is removed
     return f"{tc_url}/api/queue/v1/task/{task_id}/artifacts/public/{script}"
 

src/taskgraph/util/parameterization.py

@@ -90,11 +90,7 @@ def repl(match):
                         f"task '{label}' has no dependency named '{dependency}'"
                     )
 
-            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 get_artifact_url(task_id, artifact_name)
 
         return ARTIFACT_REFERENCE_PATTERN.sub(repl, val)