|
1 | 1 | import logging |
2 | | - |
| 2 | +import contextvars |
3 | 3 | from rich.console import Console |
4 | 4 | from rich.logging import RichHandler |
| 5 | +import json_log_formatter |
| 6 | +import os |
| 7 | +import ddtrace |
| 8 | +from ddtrace import tracer |
| 9 | + |
| 10 | +_is_datadog_configured = bool(os.environ.get("DD_AGENT_HOST")) |
| 11 | + |
| 12 | +ctx_var_request_id = contextvars.ContextVar[str]("request_id") |
| 13 | + |
| 14 | + |
| 15 | +class CustomJSONFormatter(json_log_formatter.JSONFormatter): |
| 16 | + def json_record(self, message: str, extra: dict, record: logging.LogRecord) -> dict: |
| 17 | + extra = super().json_record(message, extra, record) |
| 18 | + extra["level"] = record.levelname |
| 19 | + extra["name"] = record.name |
| 20 | + extra["lineno"] = record.lineno |
| 21 | + extra["pathname"] = record.pathname |
| 22 | + extra["request_id"] = ctx_var_request_id.get(None) |
| 23 | + if _is_datadog_configured: |
| 24 | + extra["dd.trace_id"] = tracer.get_log_correlation_context().get("dd.trace_id", None) or getattr( |
| 25 | + record, "dd.trace_id", 0 |
| 26 | + ) |
| 27 | + extra["dd.span_id"] = tracer.get_log_correlation_context().get("dd.span_id", None) or getattr( |
| 28 | + record, "dd.span_id", 0 |
| 29 | + ) |
| 30 | + # add the env, service, and version configured for the tracer |
| 31 | + # If tracing is not set up, then this should pull values from DD_ENV, DD_SERVICE, and DD_VERSION. |
| 32 | + service_override = ddtrace.config.service or os.getenv("DD_SERVICE") |
| 33 | + if service_override: |
| 34 | + extra["dd.service"] = service_override |
| 35 | + |
| 36 | + env_override = ddtrace.config.env or os.getenv("DD_ENV") |
| 37 | + if env_override: |
| 38 | + extra["dd.env"] = env_override |
5 | 39 |
|
| 40 | + version_override = ddtrace.config.version or os.getenv("DD_VERSION") |
| 41 | + if version_override: |
| 42 | + extra["dd.version"] = version_override |
6 | 43 |
|
7 | | -def make_logger(name: str): |
| 44 | + return extra |
| 45 | + |
| 46 | +def make_logger(name: str) -> logging.Logger: |
8 | 47 | """ |
9 | 48 | Creates a logger object with a RichHandler to print colored text. |
10 | 49 | :param name: The name of the module to create the logger for. |
11 | 50 | :return: A logger object. |
12 | 51 | """ |
13 | 52 | # Create a console object to print colored text |
14 | | - console = Console() |
15 | | - |
16 | | - # Create a logger object with the name of the current module |
17 | 53 | logger = logging.getLogger(name) |
18 | | - |
19 | | - # Set the global log level to INFO |
20 | 54 | logger.setLevel(logging.INFO) |
21 | 55 |
|
22 | | - # Add the RichHandler to the logger to print colored text |
23 | | - handler = RichHandler( |
24 | | - console=console, |
25 | | - show_level=False, |
26 | | - show_path=False, |
27 | | - show_time=False, |
28 | | - ) |
29 | | - logger.addHandler(handler) |
| 56 | + environment = os.getenv("ENVIRONMENT") |
| 57 | + if environment == "local": |
| 58 | + console = Console() |
| 59 | + # Add the RichHandler to the logger to print colored text |
| 60 | + handler = RichHandler( |
| 61 | + console=console, |
| 62 | + show_level=False, |
| 63 | + show_path=False, |
| 64 | + show_time=False, |
| 65 | + ) |
| 66 | + logger.addHandler(handler) |
| 67 | + return logger |
| 68 | + |
| 69 | + stream_handler = logging.StreamHandler() |
| 70 | + if _is_datadog_configured: |
| 71 | + stream_handler.setFormatter(CustomJSONFormatter()) |
| 72 | + else: |
| 73 | + stream_handler.setFormatter( |
| 74 | + logging.Formatter("%(asctime)s %(levelname)s [%(name)s] [%(filename)s:%(lineno)d] - %(message)s") |
| 75 | + ) |
30 | 76 |
|
| 77 | + logger.addHandler(stream_handler) |
| 78 | + # Create a logger object with the name of the current module |
31 | 79 | return logger |
0 commit comments