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

Enhance logging in bitbucket_app, gitlab_webhook, and google_cloud_storage_secret_provider with JSON format and additional context

Browse Source
This commit is contained in:
Ori Kotek
2023-10-18 16:44:03 +03:00
parent 88ca501c0c
commit 1314898cbf
3 changed files with 36 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
pr_agent/secret_providers/google_cloud_storage_secret_provider.py
View File

@ -1,9 +1,8 @@
import ujson import ujson
from google.cloud import storage from google.cloud import storage
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
from pr_agent.git_providers.gitlab_provider import logger from pr_agent.log import get_logger
from pr_agent.secret_providers.secret_provider import SecretProvider from pr_agent.secret_providers.secret_provider import SecretProvider
@ -15,7 +14,7 @@ class GoogleCloudStorageSecretProvider(SecretProvider):
self.bucket_name = get_settings().google_cloud_storage.bucket_name self.bucket_name = get_settings().google_cloud_storage.bucket_name
self.bucket = self.client.bucket(self.bucket_name) self.bucket = self.client.bucket(self.bucket_name)
except Exception as e: except Exception as e:
logger.error(f"Failed to initialize Google Cloud Storage Secret Provider: {e}") get_logger().error(f"Failed to initialize Google Cloud Storage Secret Provider: {e}")
raise e raise e
def get_secret(self, secret_name: str) -> str: def get_secret(self, secret_name: str) -> str:
@ -23,7 +22,7 @@ class GoogleCloudStorageSecretProvider(SecretProvider):
blob = self.bucket.blob(secret_name) blob = self.bucket.blob(secret_name)
return blob.download_as_string() return blob.download_as_string()
except Exception as e: except Exception as e:
logger.error(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}") get_logger().error(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}")
return "" return ""
def store_secret(self, secret_name: str, secret_value: str): def store_secret(self, secret_name: str, secret_value: str):
@ -31,5 +30,5 @@ class GoogleCloudStorageSecretProvider(SecretProvider):
blob = self.bucket.blob(secret_name) blob = self.bucket.blob(secret_name)
blob.upload_from_string(secret_value) blob.upload_from_string(secret_value)
except Exception as e: except Exception as e:
logger.error(f"Failed to store secret {secret_name} in Google Cloud Storage: {e}") get_logger().error(f"Failed to store secret {secret_name} in Google Cloud Storage: {e}")
raise e raise e

25
pr_agent/servers/bitbucket_app.py
View File

@ -65,15 +65,17 @@ async def handle_manifest(request: Request, response: Response):
@router.post("/webhook") @router.post("/webhook")
async def handle_github_webhooks(background_tasks: BackgroundTasks, request: Request): async def handle_github_webhooks(background_tasks: BackgroundTasks, request: Request):
print(request.headers) log_context = {"server_type": "bitbucket_app"}
get_logger().debug(request.headers)
jwt_header = request.headers.get("authorization", None) jwt_header = request.headers.get("authorization", None)
if jwt_header: if jwt_header:
input_jwt = jwt_header.split(" ")[1] input_jwt = jwt_header.split(" ")[1]
data = await request.json() data = await request.json()
print(data) get_logger().debug(data)
async def inner(): async def inner():
try: try:
owner = data["data"]["repository"]["owner"]["username"] owner = data["data"]["repository"]["owner"]["username"]
log_context["sender"] = owner
secrets = json.loads(secret_provider.get_secret(owner)) secrets = json.loads(secret_provider.get_secret(owner))
shared_secret = secrets["shared_secret"] shared_secret = secrets["shared_secret"]
client_key = secrets["client_key"] client_key = secrets["client_key"]
@ -85,11 +87,17 @@ async def handle_github_webhooks(background_tasks: BackgroundTasks, request: Req
agent = PRAgent() agent = PRAgent()
if event == "pullrequest:created": if event == "pullrequest:created":
pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] pr_url = data["data"]["pullrequest"]["links"]["html"]["href"]
await agent.handle_request(pr_url, "review") log_context["api_url"] = pr_url
log_context["event"] = "pull_request"
with get_logger().contextualize(**log_context):
await agent.handle_request(pr_url, "review")
elif event == "pullrequest:comment_created": elif event == "pullrequest:comment_created":
pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] pr_url = data["data"]["pullrequest"]["links"]["html"]["href"]
log_context["api_url"] = pr_url
log_context["event"] = "comment"
comment_body = data["data"]["comment"]["content"]["raw"] comment_body = data["data"]["comment"]["content"]["raw"]
await agent.handle_request(pr_url, comment_body) with get_logger().contextualize(**log_context):
await agent.handle_request(pr_url, comment_body)
except Exception as e: except Exception as e:
get_logger().error(f"Failed to handle webhook: {e}") get_logger().error(f"Failed to handle webhook: {e}")
background_tasks.add_task(inner) background_tasks.add_task(inner)
@ -102,9 +110,10 @@ async def handle_github_webhooks(request: Request, response: Response):
@router.post("/installed") @router.post("/installed")
async def handle_installed_webhooks(request: Request, response: Response): async def handle_installed_webhooks(request: Request, response: Response):
try: try:
print(request.headers) get_logger().info("handle_installed_webhooks")
get_logger().info(request.headers)
data = await request.json() data = await request.json()
print(data) get_logger().info(data)
shared_secret = data["sharedSecret"] shared_secret = data["sharedSecret"]
client_key = data["clientKey"] client_key = data["clientKey"]
username = data["principal"]["username"] username = data["principal"]["username"]
@ -119,8 +128,10 @@ async def handle_installed_webhooks(request: Request, response: Response):
@router.post("/uninstalled") @router.post("/uninstalled")
async def handle_uninstalled_webhooks(request: Request, response: Response): async def handle_uninstalled_webhooks(request: Request, response: Response):
get_logger().info("handle_uninstalled_webhooks")
data = await request.json() data = await request.json()
print(data) get_logger().info(data)
def start(): def start():

18
pr_agent/servers/gitlab_webhook.py
View File

@ -12,23 +12,33 @@ from starlette_context.middleware import RawContextMiddleware
from pr_agent.agent.pr_agent import PRAgent from pr_agent.agent.pr_agent import PRAgent
from pr_agent.config_loader import get_settings, global_settings from pr_agent.config_loader import get_settings, global_settings
from pr_agent.log import get_logger, setup_logger from pr_agent.log import LoggingFormat, get_logger, setup_logger
from pr_agent.secret_providers import get_secret_provider from pr_agent.secret_providers import get_secret_provider
setup_logger() setup_logger(fmt=LoggingFormat.JSON)
router = APIRouter() router = APIRouter()
secret_provider = get_secret_provider() if get_settings().get("CONFIG.SECRET_PROVIDER") else None 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):
log_context["action"] = body
log_context["event"] = "pull_request" if body == "/review" else "comment"
log_context["api_url"] = url
with get_logger().contextualize(**log_context):
background_tasks.add_task(PRAgent().handle_request, url, body)
@router.post("/webhook") @router.post("/webhook")
async def gitlab_webhook(background_tasks: BackgroundTasks, request: Request): async def gitlab_webhook(background_tasks: BackgroundTasks, request: Request):
log_context = {"server_type": "gitlab_app"}
if request.headers.get("X-Gitlab-Token") and secret_provider: if request.headers.get("X-Gitlab-Token") and secret_provider:
request_token = request.headers.get("X-Gitlab-Token") request_token = request.headers.get("X-Gitlab-Token")
secret = secret_provider.get_secret(request_token) secret = secret_provider.get_secret(request_token)
try: try:
secret_dict = json.loads(secret) secret_dict = json.loads(secret)
gitlab_token = secret_dict["gitlab_token"] gitlab_token = secret_dict["gitlab_token"]
log_context["sender"] = secret_dict["id"]
context["settings"] = copy.deepcopy(global_settings) context["settings"] = copy.deepcopy(global_settings)
context["settings"].gitlab.personal_access_token = gitlab_token context["settings"].gitlab.personal_access_token = gitlab_token
except Exception as e: except Exception as e:
@ -48,13 +58,13 @@ async def gitlab_webhook(background_tasks: BackgroundTasks, request: Request):
if data.get('object_kind') == 'merge_request' and data['object_attributes'].get('action') in ['open', 'reopen']: if data.get('object_kind') == 'merge_request' and data['object_attributes'].get('action') in ['open', 'reopen']:
get_logger().info(f"A merge request has been opened: {data['object_attributes'].get('title')}") get_logger().info(f"A merge request has been opened: {data['object_attributes'].get('title')}")
url = data['object_attributes'].get('url') url = data['object_attributes'].get('url')
background_tasks.add_task(PRAgent().handle_request, url, "/review") handle_request(background_tasks, url, "/review")
elif data.get('object_kind') == 'note' and data['event_type'] == 'note': elif data.get('object_kind') == 'note' and data['event_type'] == 'note':
if 'merge_request' in data: if 'merge_request' in data:
mr = data['merge_request'] mr = data['merge_request']
url = mr.get('url') url = mr.get('url')
body = data.get('object_attributes', {}).get('note') body = data.get('object_attributes', {}).get('note')
background_tasks.add_task(PRAgent().handle_request, url, body) handle_request(background_tasks, url, body)
return JSONResponse(status_code=status.HTTP_200_OK, content=jsonable_encoder({"message": "success"})) return JSONResponse(status_code=status.HTTP_200_OK, content=jsonable_encoder({"message": "success"}))