From 413547f404f9b4023880d36d3516a30d60929a28 Mon Sep 17 00:00:00 2001 From: Ori Kotek Date: Mon, 3 Jun 2024 12:36:24 +0300 Subject: [PATCH] Refactor GitLab webhook handling for async processing and improved logging --- pr_agent/servers/gitlab_webhook.py | 101 ++++++++++++++++------------- 1 file changed, 55 insertions(+), 46 deletions(-) diff --git a/pr_agent/servers/gitlab_webhook.py b/pr_agent/servers/gitlab_webhook.py index 715185a8..3ee4a08c 100644 --- a/pr_agent/servers/gitlab_webhook.py +++ b/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"} - get_logger().debug("Received a GitLab webhook") + start_time = datetime.now() + request_json = await request.json() - # Check if the request is authorized - if request.headers.get("X-Gitlab-Token") and secret_provider: - request_token = request.headers.get("X-Gitlab-Token") - secret = secret_provider.get_secret(request_token) - try: - secret_dict = json.loads(secret) - gitlab_token = secret_dict["gitlab_token"] - log_context["sender"] = secret_dict.get("token_name", secret_dict.get("id", "unknown")) - context["settings"] = copy.deepcopy(global_settings) - context["settings"].gitlab.personal_access_token = gitlab_token - except Exception as e: - get_logger().error(f"Failed to validate secret {request_token}: {e}") + async def inner(data: dict): + log_context = {"server_type": "gitlab_app"} + get_logger().debug("Received a GitLab webhook") + if request.headers.get("X-Gitlab-Token") and secret_provider: + request_token = request.headers.get("X-Gitlab-Token") + secret = secret_provider.get_secret(request_token) + try: + secret_dict = json.loads(secret) + gitlab_token = secret_dict["gitlab_token"] + log_context["token_id"] = secret_dict.get("token_name", secret_dict.get("id", "unknown")) + context["settings"] = copy.deepcopy(global_settings) + 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"): - secret = get_settings().get("GITLAB.SHARED_SECRET") - if not request.headers.get("X-Gitlab-Token") == secret: - get_logger().error(f"Failed to validate secret") + gitlab_token = get_settings().get("GITLAB.PERSONAL_ACCESS_TOKEN", None) + if not gitlab_token: + get_logger().error("No gitlab token found") 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']: - 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) + await handle_request(url, body, log_context, sender_id) + 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}")