refactor(api): move signal receiver out from init_app

This ensures that we can write unit tests for them.
pull/19783/head
QuantumGhost 1 year ago
parent b0051dce39
commit 4bd57976d7

@ -2,60 +2,71 @@ import json
import logging
import flask
import werkzeug.http
from flask import Flask
from flask.signals import request_finished, request_started
from configs import dify_config
logger = logging.getLogger(__name__)
_logger = logging.getLogger(__name__)
def _is_content_type_json(content_type: str) -> bool:
if not content_type:
return False
content_type_no_option, _ = werkzeug.http.parse_options_header(content_type)
return content_type_no_option.lower() == "application/json"
def _log_request_started(_sender, **_extra):
"""Log the start of a request."""
if not _logger.isEnabledFor(logging.DEBUG):
return
request = flask.request
if not (_is_content_type_json(request.content_type) and request.data):
_logger.debug("Received Request %s -> %s", request.method, request.path)
return
try:
json_data = json.loads(request.data)
except (TypeError, ValueError):
_logger.exception("Failed to parse JSON request")
formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
_logger.debug(
"Received Request %s -> %s, Request Body:\n%s",
request.method,
request.path,
formatted_json,
)
def _log_request_finished(_sender, response, **_extra):
"""Log the end of a request."""
if not _logger.isEnabledFor(logging.DEBUG) or response is None:
return
if not _is_content_type_json(response.content_type):
_logger.debug("Response %s %s", response.status, response.content_type)
return
response_data = response.get_data(as_text=True)
try:
json_data = json.loads(response_data)
except (TypeError, ValueError):
_logger.exception("Failed to parse JSON response")
formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
_logger.debug(
"Response %s %s, Response Body:\n%s",
response.status,
response.content_type,
formatted_json,
)
def init_app(app: Flask):
"""Initialize the request logging extension."""
if not dify_config.ENABLE_REQUEST_LOGGING:
return
@request_started.connect_via(app)
def log_request_started(sender, **extra):
"""Log the start of a request."""
if not logger.isEnabledFor(logging.DEBUG):
return
request = flask.request
content_type = request.content_type
if request.data and "application/json" in content_type.lower():
try:
json_data = json.loads(request.data)
formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
logger.debug(
"Received Request %s -> %s, Request Body:\n%s",
request.method,
request.path,
formatted_json,
)
except Exception:
logger.exception("Failed to parse JSON request")
else:
logger.debug("Received Request %s -> %s", request.method, request.path)
@request_finished.connect_via(app)
def log_request_finished(sender, response, **extra):
"""Log the end of a request."""
if not logger.isEnabledFor(logging.DEBUG) or response is None:
return
if response.content_type and "application/json" in response.content_type.lower():
try:
response_data = response.get_data(as_text=True)
json_data = json.loads(response_data)
formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
logger.debug(
"Response %s %s, Response Body:\n%s",
response.status,
response.content_type,
formatted_json,
)
except Exception:
logger.exception("Failed to parse JSON response")
else:
logger.debug("Response %s %s", response.status, response.content_type)
request_started.connect(_log_request_started, app)
request_finished.connect(_log_request_finished, app)

Loading…
Cancel
Save