Pull some info on the role for batch tasks

This might be helpful for debugging

Author: namelessjon

airline/awsbatch.py

@@ -1,10 +1,14 @@
 import functools
+import json
 import os
+import urllib.request as req
+import time
+import random
 
 import airline
 
 
-def airline_wrapper(_handler=None, *, env_vars=None):
+def airline_wrapper(_handler=None, *, env_vars=None, add_role_info=True):
     '''airline decorator for a function in a AWS Batch job
     ```
     @airline_wrapper
@@ -35,6 +39,9 @@ def _airline_wrapper(*args, **kwargs):
                     for name, var in env_vars.items():
                         airline.add_environment_variable(name, var)
 
+                if add_role_info:
+                    _add_role_info()
+
                 resp = handler(*args, **kwargs)
 
                 return resp
@@ -44,3 +51,35 @@ def _airline_wrapper(*args, **kwargs):
         return decorator_airline
     else:
         return decorator_airline(_handler)
+
+
+def _add_role_info():
+    uri = os.getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")
+    if not uri:
+        return
+
+    full_uri = f"169.254.170.2{uri}"
+    data = _fetch_parse_uri(full_uri)
+
+    airline.add_context_field('aws.role_arn', data.get('RoleArn'))
+    airline.add_context_field('aws.access_key_id', data.get('AccessKeyId'))
+
+
+def _fetch_parse_uri(uri, attempts=5):
+    attempt = 0
+    while True:
+        try:
+            attempt += 1
+            with req.urlopen(uri, timeout=1) as f:
+                response = json.load(f)
+                return response
+        except req.URLError:
+            if attempt < attempts:
+                _sleep(attempt)
+            else:
+                raise
+
+
+def _sleep(attempt, cap=60, base=1):
+    sleep_time = random.uniform(0, min(cap, base * 2**attempt))
+    time.sleep(sleep_time)