Inject taskid in job submission result (#72)

* Inject taskid in job submission result

* Fix formatting

firecrest/AsyncClient.py

@@ -17,15 +17,7 @@
 import time
 
 from contextlib import nullcontext
-from typing import (
-    Any,
-    ContextManager,
-    Optional,
-    overload,
-    Sequence,
-    Tuple,
-    List,
-)
+from typing import Any, ContextManager, Optional, overload, Sequence, Tuple, List
 from requests.compat import json  # type: ignore
 from packaging.version import Version, parse
 
@@ -116,7 +108,7 @@ async def wrapper(*args, **kwargs):
                         default=resp.headers.get("RateLimit-Reset", default=10),
                     )
                     reset = int(reset)
-                    microservice = kwargs['endpoint'].split("/")[1]
+                    microservice = kwargs["endpoint"].split("/")[1]
                     client = args[0]
                     logger.info(
                         f"Rate limit in `{microservice}` is reached, next "
@@ -194,7 +186,9 @@ def set_api_version(self, api_version: str) -> None:
         self._api_version = parse(api_version)
 
     @_retry_requests  # type: ignore
-    async def _get_request(self, endpoint, additional_headers=None, params=None) -> httpx.Response:
+    async def _get_request(
+        self, endpoint, additional_headers=None, params=None
+    ) -> httpx.Response:
         microservice = endpoint.split("/")[1]
         url = f"{self._firecrest_url}{endpoint}"
         async with self._locks[microservice]:
@@ -240,7 +234,9 @@ async def _post_request(
         return resp
 
     @_retry_requests  # type: ignore
-    async def _put_request(self, endpoint, additional_headers=None, data=None) -> httpx.Response:
+    async def _put_request(
+        self, endpoint, additional_headers=None, data=None
+    ) -> httpx.Response:
         microservice = endpoint.split("/")[1]
         url = f"{self._firecrest_url}{endpoint}"
         async with self._locks[microservice]:
@@ -980,7 +976,11 @@ async def submit(
         json_response = self._json_response([resp], 201)
         logger.info(f"Job submission task: {json_response['task_id']}")
         t = ComputeTask(self, json_response["task_id"], [resp])
-        return await t.poll_task("200")
+
+        # Inject taskid in the result
+        result = await t.poll_task("200")
+        result["firecrest_taskid"] = json_response["task_id"]
+        return result
 
     async def poll(
         self,

firecrest/AsyncExternalStorage.py

@@ -14,12 +14,7 @@
 import requests
 import shutil
 import sys
-from typing import (
-    ContextManager,
-    Optional,
-    List,
-    TYPE_CHECKING,
-)
+from typing import ContextManager, Optional, List, TYPE_CHECKING
 import urllib.request
 from packaging.version import Version
 

firecrest/BasicClient.py

@@ -859,9 +859,13 @@ def submit(
         self._current_method_requests = []
         json_response = self._submit_request(machine, job_script, local_file, account)
         logger.info(f"Job submission task: {json_response['task_id']}")
-        return self._poll_tasks(
+
+        # Inject taskid in the result
+        result = self._poll_tasks(
             json_response["task_id"], "200", itertools.cycle([1, 5, 10])
         )
+        result["firecrest_taskid"] = json_response["task_id"]
+        return result
 
     def poll(
         self,