Coverage for polar/middlewares.py: 76%

1import functools 

2import re 

3 

4import dramatiq 

5import structlog 

6from starlette.datastructures import MutableHeaders 

7from starlette.types import ASGIApp, Message, Receive, Scope, Send 

8 

9from polar.logging import Logger, generate_correlation_id 

10from polar.worker import JobQueueManager 

11 

12 

13class LogCorrelationIdMiddleware: 

14 def __init__(self, app: ASGIApp) -> None: 

15 self.app = app 

16 

17 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

18 if scope["type"] != "http": 

19 return await self.app(scope, receive, send) 

20 

21 structlog.contextvars.bind_contextvars( 

22 correlation_id=generate_correlation_id(), 

23 method=scope["method"], 

24 path=scope["path"], 

25 ) 

26 

27 await self.app(scope, receive, send) 

28 

29 structlog.contextvars.unbind_contextvars("correlation_id", "method", "path") 

30 

31 

32class FlushEnqueuedWorkerJobsMiddleware: 

33 def __init__(self, app: ASGIApp) -> None: 

34 self.app = app 

35 

36 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

37 if scope["type"] not in ("http", "websocket"): 

38 await self.app(scope, receive, send) 

39 return 

40 

41 async with JobQueueManager.open(dramatiq.get_broker(), scope["state"]["redis"]): 

42 await self.app(scope, receive, send) 

43 

44 

45class PathRewriteMiddleware: 

46 def __init__( 

47 self, app: ASGIApp, pattern: str | re.Pattern[str], replacement: str 

48 ) -> None: 

49 self.app = app 

50 self.pattern = pattern 

51 self.replacement = replacement 

52 self.logger: Logger = structlog.get_logger() 

53 

54 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

55 if scope["type"] not in ("http", "websocket"): 

56 await self.app(scope, receive, send) 

57 return 

58 

59 scope["path"], replacements = re.subn( 

60 self.pattern, self.replacement, scope["path"] 

61 ) 

62 

63 if replacements > 0: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true

64 self.logger.warning( 

65 "PathRewriteMiddleware", 

66 pattern=self.pattern, 

67 replacement=self.replacement, 

68 path=scope["path"], 

69 ) 

70 

71 send = functools.partial(self.send, send=send, replacements=replacements) 

72 await self.app(scope, receive, send) 

73 

74 async def send(self, message: Message, send: Send, replacements: int) -> None: 

75 if message["type"] != "http.response.start": 

76 await send(message) 

77 return 

78 

79 message.setdefault("headers", []) 

80 headers = MutableHeaders(scope=message) 

81 if replacements > 0: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 headers["X-Polar-Deprecation-Notice"] = ( 

83 "The API root has moved from /api/v1 to /v1. " 

84 "Please update your integration." 

85 ) 

86 

87 await send(message) 

88 

89 

90class SandboxResponseHeaderMiddleware: 

91 def __init__(self, app: ASGIApp) -> None: 

92 self.app = app 

93 

94 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

95 if scope["type"] not in ("http", "websocket"): 

96 await self.app(scope, receive, send) 

97 return 

98 

99 async def send_wrapper(message: Message) -> None: 

100 if message["type"] == "http.response.start": 

101 message.setdefault("headers", []) 

102 headers = MutableHeaders(scope=message) 

103 headers["X-Polar-Sandbox"] = "1" 

104 await send(message) 

105 

106 await self.app(scope, receive, send_wrapper)