Repository: Netflix/consoleme
File: consoleme/handlers/base.py (Python)

Code before:
    @staticmethod
    async def init_saml_auth(req):
        auth = await sync_to_async(OneLogin_Saml2_Auth)(
            req, custom_base_path=config.get("get_user_by_saml_settings.saml_path")
        )
        return auth


class BaseAPIV1Handler(BaseHandler):
    """Default API Handler for api/v1/* routes."""
    def set_default_headers(self) -> None:
        self.set_header("Content-Type", "application/json")


class BaseAPIV2Handler(BaseHandler):
    """Default API Handler for api/v2/* routes."""
    def set_default_headers(self) -> None:
        self.set_header("Content-Type", "application/json")

    def write_error(self, status_code: int, **kwargs: Any) -> None:
        if self.settings.get("serve_traceback") and "exc_info" in kwargs:
            # in debug mode, try to send a traceback
            self.set_header("Content-Type", "text/plain")
            for line in traceback.format_exception(*kwargs["exc_info"]):
                self.write(line)
            self.finish()
        else:
            self.set_header("Content-Type", "application/problem+json")
            title = httputil.responses.get(status_code, "Unknown")
            message = kwargs.get("message", self._reason)
            # self.set_status() modifies self._reason, so this call should come after we grab the reason
            self.set_status(status_code)
            self.finish(
                json.dumps(
                    {"status": status_code, "title": title, "message": message}
                )  # noqa
            )


class BaseMtlsHandler(BaseAPIV1Handler):
    def initialize(self, **kwargs):
        self.kwargs = kwargs

    async def prepare(self):
        self.tracer = None
        self.span = None
        self.spans = {}
        self.responses = []
        self.request_uuid = str(uuid.uuid4())
        stats.timer("base_handler.incoming_request")
        try:

Review comment (suggestion):
Instead of writing this as the response, what would you think of adding it as an element in the response JSON? That way we maintain the status code and the structure of the error response and also have access to the traceback.

Code after:
    @staticmethod
    async def init_saml_auth(req):
        auth = await sync_to_async(OneLogin_Saml2_Auth)(
            req, custom_base_path=config.get("get_user_by_saml_settings.saml_path")
        )
        return auth


class BaseAPIV1Handler(BaseHandler):
    """Default API Handler for api/v1/* routes."""
    def set_default_headers(self) -> None:
        self.set_header("Content-Type", "application/json")


class BaseAPIV2Handler(BaseHandler):
    """Default API Handler for api/v2/* routes."""
    def set_default_headers(self) -> None:
        self.set_header("Content-Type", "application/json")

    def write_error(self, status_code: int, **kwargs: Any) -> None:
        if self.settings.get("serve_traceback") and "exc_info" in kwargs:
            # in debug mode, try to send a traceback
            self.set_header("Content-Type", "text/plain")
            self.set_status(status_code)
            for line in traceback.format_exception(*kwargs["exc_info"]):
                self.write(line)
            self.finish()
        else:
            self.set_header("Content-Type", "application/problem+json")
            title = httputil.responses.get(status_code, "Unknown")
            message = kwargs.get("message", self._reason)
            # self.set_status() modifies self._reason, so this call should come after we grab the reason
            self.set_status(status_code)
            self.finish(
                json.dumps(
                    {"status": status_code, "title": title, "message": message}
                )  # noqa
            )


class BaseMtlsHandler(BaseAPIV1Handler):
    def initialize(self, **kwargs):
        self.kwargs = kwargs

    async def prepare(self):
        self.tracer = None
        self.span = None
        self.spans = {}
        self.responses = []
        self.request_uuid = str(uuid.uuid4())
        stats.timer("base_handler.incoming_request")