Coverage for src/local_deep_research/utilities/threading_utils.py: 100%

1import threading 

2import uuid 

3from functools import wraps 

4from typing import Any, Callable, Hashable, Tuple 

5 

6from cachetools import cached, keys 

7from flask import current_app, g 

8from flask.ctx import AppContext 

9from loguru import logger 

10 

11g_thread_local_store = threading.local() 

12 

13# Keys in Flask ``g`` that are created and cleaned up by a 

14# ``teardown_appcontext`` handler (see ``web/app_factory.cleanup_db_session``). 

15# They must never be copied into a worker context: the value is a resource 

16# (a SQLAlchemy session) still owned by the submitting thread, and copying it 

17# would let a teardown roll back and close the parent thread's session. 

18_TEARDOWN_OWNED_G_KEYS = frozenset({"db_session"}) 

19 

20 

21def thread_specific_cache(*args: Any, **kwargs: Any) -> Callable: 

22 """ 

23 A version of `cached()` that is local to a single thread. In other words, 

24 cache entries will only be valid in the thread where they were created. 

25 

26 Args: 

27 *args: Will be forwarded to `cached()`. 

28 **kwargs: Will be forwarded to `cached()`. 

29 

30 Returns: 

31 The wrapped function. 

32 

33 """ 

34 

35 def _key_func(*args_: Any, **kwargs_: Any) -> Tuple[Hashable, ...]: 

36 base_hash = keys.hashkey(*args_, **kwargs_) 

37 

38 if hasattr(g_thread_local_store, "thread_id"): 

39 # We already gave this thread a unique ID. Use that. 

40 thread_id = g_thread_local_store.thread_id 

41 else: 

42 # Give this thread a new unique ID. 

43 thread_id = uuid.uuid4().hex 

44 g_thread_local_store.thread_id = thread_id 

45 

46 return (thread_id,) + base_hash 

47 

48 return cached(*args, **kwargs, key=_key_func) 

49 

50 

51def thread_with_app_context(to_wrap: Callable) -> Callable: 

52 """ 

53 Decorator that wraps the entry point to a thread and injects the current 

54 app context from Flask. This is useful when we want to use multiple 

55 threads to handle a single request. 

56 

57 When using this wrapped function, `current_app.app_context()` should be 

58 passed as the first argument when initializing the thread. 

59 

60 Args: 

61 to_wrap: The function to wrap. 

62 

63 Returns: 

64 The wrapped function. 

65 

66 """ 

67 

68 @wraps(to_wrap) 

69 def _run_with_context( 

70 app_context: AppContext | None, *args: Any, **kwargs: Any 

71 ) -> Any: 

72 if app_context is None: 

73 # Do nothing. 

74 return to_wrap(*args, **kwargs) 

75 

76 with app_context: 

77 return to_wrap(*args, **kwargs) 

78 

79 return _run_with_context 

80 

81 

82def thread_context() -> AppContext | None: 

83 """ 

84 Builds (but does not push) a new app context for a thread that is being 

85 spawned to handle the current request, pre-populating its ``g`` with the 

86 current context's global data. Teardown-owned resources such as 

87 ``db_session`` are not copied, so the worker never shares the submitting 

88 thread's session. The caller is responsible for entering the returned 

89 context in the worker thread (see :func:`thread_with_app_context`). 

90 

91 Returns: 

92 The new context, or None if no context is active. 

93 

94 """ 

95 # Copy global data. 

96 global_data = {} 

97 try: 

98 for key in g: 

99 global_data[key] = g.get(key) 

100 except (TypeError, RuntimeError): 

101 # Context is not initialized (TypeError) or no request/app context 

102 # is active (RuntimeError). Don't change anything. 

103 logger.debug( 

104 "No active request context, skipping Flask g copy for child thread" 

105 ) 

106 

107 try: 

108 context = current_app.app_context() 

109 except RuntimeError: 

110 # Context is not initialized. 

111 logger.debug("No current app context, not passing to thread.") 

112 return None 

113 

114 # Populate the new context's ``g`` directly instead of pushing the context 

115 # on the submitting thread. Pushing and popping it here would fire the app's 

116 # ``teardown_appcontext`` handlers on this thread, and a handler that owns 

117 # ``db_session`` (``web/app_factory.cleanup_db_session``) rolls back and 

118 # closes it. Since ``g`` is shallow-copied, that session is the one the 

119 # submitting thread still holds, so the parent's uncommitted work would be 

120 # discarded. Teardown-owned keys are excluded for the same reason: a worker 

121 # that needs the database acquires its own context-local session (see 

122 # ``database/session_context.get_user_db_session``) rather than reusing the 

123 # parent's. The worker enters ``with context:`` itself (via 

124 # ``thread_with_app_context``), so its own teardown cleans up only what it 

125 # created. 

126 for key, value in global_data.items(): 

127 if key in _TEARDOWN_OWNED_G_KEYS: 

128 continue 

129 setattr(context.g, key, value) 

130 

131 return context