1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
|
import uuid
from typing import TYPE_CHECKING, Any, List, Mapping, Tuple, Type
import structlog
from django_extensions.management.signals import ( # type: ignore[import-untyped]
post_command,
pre_command,
)
if TYPE_CHECKING: # pragma: no cover
import contextvars
logger = structlog.getLogger(__name__)
class DjangoCommandReceiver:
stack: List[Tuple[str, Mapping[str, "contextvars.Token[Any]"]]]
def __init__(self) -> None:
self.stack = []
def pre_receiver(self, sender: Type[Any], *args: Any, **kwargs: Any) -> None:
command_id = str(uuid.uuid4())
if len(self.stack):
parent_command_id, _ = self.stack[-1]
tokens = structlog.contextvars.bind_contextvars(
parent_command_id=parent_command_id, command_id=command_id
)
else:
tokens = structlog.contextvars.bind_contextvars(command_id=command_id)
self.stack.append((command_id, tokens))
logger.info(
"command_started",
command_name=sender.__module__.replace(".management.commands", ""),
)
def post_receiver(
self, sender: Type[Any], outcome: str, *args: Any, **kwargs: Any
) -> None:
logger.info("command_finished")
if len(self.stack): # pragma: no branch
command_id, tokens = self.stack.pop()
structlog.contextvars.reset_contextvars(**tokens)
def connect_signals(self) -> None:
pre_command.connect(self.pre_receiver)
post_command.connect(self.post_receiver)
|
