Refactor GitLab webhook handling for async processing and improved logging

pr_agent/servers/gitlab_webhook.py

@@ -1,5 +1,6 @@
 import copy
 import json
+from datetime import datetime
 
 import uvicorn
 from fastapi import APIRouter, FastAPI, Request, status
@@ -23,15 +24,17 @@ router = APIRouter()
 secret_provider = get_secret_provider() if get_settings().get("CONFIG.SECRET_PROVIDER") else None
 
 
-def handle_request(background_tasks: BackgroundTasks, url: str, body: str, log_context: dict):
+async def handle_request(api_url: str, body: str, log_context: dict, sender_id: str):
     log_context["action"] = body
     log_context["event"] = "pull_request" if body == "/review" else "comment"
-    log_context["api_url"] = url
+    log_context["api_url"] = api_url
+
     with get_logger().contextualize(**log_context):
-        background_tasks.add_task(PRAgent().handle_request, url, body)
+        await PRAgent().handle_request(api_url, body)
 
 
-async def _perform_commands_gitlab(commands_conf: str, agent: PRAgent, api_url: str, log_context: dict):
+async def _perform_commands_gitlab(commands_conf: str, agent: PRAgent, api_url: str,
+                                   log_context: dict):
     apply_repo_settings(api_url)
     commands = get_settings().get(f"gitlab.{commands_conf}", {})
     for command in commands:
@@ -50,53 +53,59 @@ async def _perform_commands_gitlab(commands_conf: str, agent: PRAgent, api_url:
 
 @router.post("/webhook")
 async def gitlab_webhook(background_tasks: BackgroundTasks, request: Request):
-    log_context = {"server_type": "gitlab_app"}
+    start_time = datetime.now()
-    get_logger().debug("Received a GitLab webhook")
+    request_json = await request.json()
 
-    # Check if the request is authorized
+    async def inner(data: dict):
-    if request.headers.get("X-Gitlab-Token") and secret_provider:
+        log_context = {"server_type": "gitlab_app"}
-        request_token = request.headers.get("X-Gitlab-Token")
+        get_logger().debug("Received a GitLab webhook")
-        secret = secret_provider.get_secret(request_token)
+        if request.headers.get("X-Gitlab-Token") and secret_provider:
-        try:
+            request_token = request.headers.get("X-Gitlab-Token")
-            secret_dict = json.loads(secret)
+            secret = secret_provider.get_secret(request_token)
-            gitlab_token = secret_dict["gitlab_token"]
+            try:
-            log_context["sender"] = secret_dict.get("token_name", secret_dict.get("id", "unknown"))
+                secret_dict = json.loads(secret)
-            context["settings"] = copy.deepcopy(global_settings)
+                gitlab_token = secret_dict["gitlab_token"]
-            context["settings"].gitlab.personal_access_token = gitlab_token
+                log_context["token_id"] = secret_dict.get("token_name", secret_dict.get("id", "unknown"))
-        except Exception as e:
+                context["settings"] = copy.deepcopy(global_settings)
-            get_logger().error(f"Failed to validate secret {request_token}: {e}")
+                context["settings"].gitlab.personal_access_token = gitlab_token
+            except Exception as e:
+                get_logger().error(f"Failed to validate secret {request_token}: {e}")
+                return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
+        elif get_settings().get("GITLAB.SHARED_SECRET"):
+            secret = get_settings().get("GITLAB.SHARED_SECRET")
+            if not request.headers.get("X-Gitlab-Token") == secret:
+                get_logger().error("Failed to validate secret")
+                return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
+        else:
+            get_logger().error("Failed to validate secret")
             return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
-    elif get_settings().get("GITLAB.SHARED_SECRET"):
+        gitlab_token = get_settings().get("GITLAB.PERSONAL_ACCESS_TOKEN", None)
-        secret = get_settings().get("GITLAB.SHARED_SECRET")
+        if not gitlab_token:
-        if not request.headers.get("X-Gitlab-Token") == secret:
+            get_logger().error("No gitlab token found")
-            get_logger().error(f"Failed to validate secret")
             return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
-    else:
-        get_logger().error(f"Failed to validate secret")
-        return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
-    gitlab_token = get_settings().get("GITLAB.PERSONAL_ACCESS_TOKEN", None)
-    if not gitlab_token:
-        get_logger().error(f"No gitlab token found")
-        return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content=jsonable_encoder({"message": "unauthorized"}))
 
-    data = await request.json()
+        get_logger().info("GitLab data", artifact=data)
-    get_logger().info("GitLab data", artifact=data)
+        sender = data.get("user", {}).get("username", "unknown")
+        sender_id = data.get("user", {}).get("id", "unknown")
+        log_context["sender"] = sender
+        if data.get('object_kind') == 'merge_request' and data['object_attributes'].get('action') in ['open', 'reopen']:
+            url = data['object_attributes'].get('url')
+            get_logger().info(f"New merge request: {url}")
+            await _perform_commands_gitlab("pr_commands", PRAgent(), url, log_context)
+        elif data.get('object_kind') == 'note' and data['event_type'] == 'note': # comment on MR
+            if 'merge_request' in data:
+                mr = data['merge_request']
+                url = mr.get('url')
+                get_logger().info(f"A comment has been added to a merge request: {url}")
+                body = data.get('object_attributes', {}).get('note')
+                if data.get('object_attributes', {}).get('type') == 'DiffNote' and '/ask' in body: # /ask_line
+                    body = handle_ask_line(body, data)
 
-    if data.get('object_kind') == 'merge_request' and data['object_attributes'].get('action') in ['open', 'reopen']:
+                await handle_request(url, body, log_context, sender_id)
-        url = data['object_attributes'].get('url')
-        get_logger().info(f"New merge request: {url}")
-        await _perform_commands_gitlab("pr_commands", PRAgent(), url, log_context)
-    elif data.get('object_kind') == 'note' and data['event_type'] == 'note': # comment on MR
-        if 'merge_request' in data:
-            mr = data['merge_request']
-            url = mr.get('url')
-            get_logger().info(f"A comment has been added to a merge request: {url}")
-            body = data.get('object_attributes', {}).get('note')
-            if data.get('object_attributes', {}).get('type') == 'DiffNote' and '/ask' in body: # /ask_line
-                body = handle_ask_line(body, data)
-
-            handle_request(background_tasks, url, body, log_context)
 
+    background_tasks.add_task(inner, request_json)
+    end_time = datetime.now()
+    get_logger().info(f"Processing time: {end_time - start_time}", request=request_json)
     return JSONResponse(status_code=status.HTTP_200_OK, content=jsonable_encoder({"message": "success"}))
 
 
@@ -113,7 +122,7 @@ def handle_ask_line(body, data):
         path = data['object_attributes']['position']['new_path']
         side = 'RIGHT'  # if line_range_['start']['type'] == 'new' else 'LEFT'
         comment_id = data['object_attributes']["discussion_id"]
-        get_logger().info(f"Handling line comment")
+        get_logger().info("Handling line comment")
         body = f"/ask_line --line_start={start_line} --line_end={end_line} --side={side} --file_name={path} --comment_id={comment_id} {question}"
     except Exception as e:
         get_logger().error(f"Failed to handle ask line comment: {e}")