xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-02 11:50:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor GitLab webhook handling for async processing and improved logging

Browse Source
This commit is contained in:
Ori Kotek
2024-06-03 12:36:24 +03:00
parent 78cabf28a4
commit 413547f404
1 changed files with 55 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
pr_agent/servers/gitlab_webhook.py
View File

@ -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}")