From 5574143ee73a3bb865fa742833e80130877c810f Mon Sep 17 00:00:00 2001
From: Michael Franklin <22381693+illusional@users.noreply.github.com>
Date: Wed, 3 Apr 2024 16:26:44 +1100
Subject: [PATCH] Capture specific environment variables on write (#721)

* Capture specific environment variables on write

* Make get_extra_audit_log_values more defensive

* Linting

---------

Co-authored-by: Michael Franklin <illusional@users.noreply.github.com>
---
 api/utils/db.py                       | 28 +++++++++++++++++++++++++++
 openapi-templates/api_client.mustache | 19 ++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/api/utils/db.py b/api/utils/db.py
index 5d88a2b1b..c574c1278 100644
--- a/api/utils/db.py
+++ b/api/utils/db.py
@@ -1,3 +1,4 @@
+import json
 import logging
 from os import getenv
 
@@ -29,6 +30,19 @@ def get_ar_guid(request: Request) -> str | None:
     return request.headers.get('sm-ar-guid')
 
 
+def get_extra_audit_log_values(request: Request) -> dict | None:
+    """Get a JSON encoded dictionary from the 'sm-extra-values' header if it exists"""
+    headers = request.headers.get('sm-extra-values')
+    if not headers:
+        return None
+
+    try:
+        return json.loads(headers)
+    except json.JSONDecodeError:
+        logging.error(f'Could not parse sm-extra-values: {headers}')
+        return None
+
+
 def get_on_behalf_of(request: Request) -> str | None:
     """
     Get sm-on-behalf-of if there are requests that were performed on behalf of
@@ -69,12 +83,16 @@ async def dependable_get_write_project_connection(
     request: Request,
     author: str = Depends(authenticate),
     ar_guid: str = Depends(get_ar_guid),
+    extra_values: dict | None = Depends(get_extra_audit_log_values),
     on_behalf_of: str | None = Depends(get_on_behalf_of),
 ) -> Connection:
     """FastAPI handler for getting connection WITH project"""
     meta = {'path': request.url.path}
     if request.client:
         meta['ip'] = request.client.host
+    if extra_values:
+        meta.update(extra_values)
+
     return await ProjectPermissionsTable.get_project_connection(
         project_name=project,
         author=author,
@@ -89,14 +107,20 @@ async def dependable_get_readonly_project_connection(
     project: str,
     author: str = Depends(authenticate),
     ar_guid: str = Depends(get_ar_guid),
+    extra_values: dict | None = Depends(get_extra_audit_log_values),
 ) -> Connection:
     """FastAPI handler for getting connection WITH project"""
+    meta = {}
+    if extra_values:
+        meta.update(extra_values)
+
     return await ProjectPermissionsTable.get_project_connection(
         project_name=project,
         author=author,
         readonly=True,
         on_behalf_of=None,
         ar_guid=ar_guid,
+        meta=meta,
     )
 
 
@@ -104,12 +128,16 @@ async def dependable_get_connection(
     request: Request,
     author: str = Depends(authenticate),
     ar_guid: str = Depends(get_ar_guid),
+    extra_values: dict | None = Depends(get_extra_audit_log_values),
 ):
     """FastAPI handler for getting connection withOUT project"""
     meta = {'path': request.url.path}
     if request.client:
         meta['ip'] = request.client.host
 
+    if extra_values:
+        meta.update(extra_values)
+
     return await SMConnections.get_connection_no_project(
         author, ar_guid=ar_guid, meta=meta
     )
diff --git a/openapi-templates/api_client.mustache b/openapi-templates/api_client.mustache
index d78bf1909..cf7296202 100644
--- a/openapi-templates/api_client.mustache
+++ b/openapi-templates/api_client.mustache
@@ -35,6 +35,21 @@ from {{packageName}}.model_utils import (
     validate_and_convert_types
 )
 
+def get_select_env_values():
+    env_values = {
+        'HAIL_ATTEMPT_ID': 'HAIL_ATTEMPT_ID',
+        'HAIL_BATCH_ID': 'HAIL_BATCH_ID',
+        'HAIL_JOB_ID': 'HAIL_JOB_ID',
+    }
+
+    as_map = {}
+    for env_key, dict_key in env_values.items():
+        value = os.getenv(env_key)
+        if value:
+            as_map[dict_key] = value
+
+    return as_map
+
 
 class ApiClient(object):
     """Generic API client for OpenAPI client library builds.
@@ -74,6 +89,10 @@ class ApiClient(object):
             self.default_headers['sm-ar-guid'] = ar_guid
         if header_name is not None:
             self.default_headers[header_name] = header_value
+        extra_values = get_select_env_values()
+        if extra_values:
+            self.default_headers['sm-extra-values'] = json.dumps(extra_values)
+
         self.cookie = cookie
         # Set default User-Agent.
         self.user_agent = '{{{httpUserAgent}}}{{^httpUserAgent}}OpenAPI-Generator/{{{packageVersion}}}/python{{/httpUserAgent}}'