Coverage for src/dataknobs_data/pooling/base.py: 24%

1"""Base classes for generic connection pool management.""" 

2 

3import asyncio 

4import atexit 

5import logging 

6from typing import Dict, Optional, Any, Protocol, TypeVar, Generic, Callable, Awaitable 

7from weakref import WeakValueDictionary 

8from abc import abstractmethod 

9 

10logger = logging.getLogger(__name__) 

11 

12 

13class PoolProtocol(Protocol): 

14 """Protocol for connection pools.""" 

15 

16 async def acquire(self): 

17 """Acquire a connection from the pool.""" 

18 ... 

19 

20 async def close(self): 

21 """Close the pool.""" 

22 ... 

23 

24 

25PoolType = TypeVar('PoolType', bound=PoolProtocol) 

26 

27 

28class BasePoolConfig: 

29 """Base configuration for connection pools.""" 

30 

31 @abstractmethod 

32 def to_connection_string(self) -> str: 

33 """Convert configuration to a connection string.""" 

34 ... 

35 

36 @abstractmethod 

37 def to_hash_key(self) -> tuple: 

38 """Create a hashable key for this configuration.""" 

39 ... 

40 

41 

42class ConnectionPoolManager(Generic[PoolType]): 

43 """ 

44 Generic connection pool manager that handles pools per event loop. 

45  

46 This class ensures that each event loop gets its own connection pool, 

47 preventing cross-loop usage errors that can occur with async connections. 

48  

49 Type Parameters: 

50 PoolType: The type of pool being managed (e.g., asyncpg.Pool) 

51 """ 

52 

53 def __init__(self): 

54 """Initialize the connection pool manager.""" 

55 # Map of (config_hash, loop_id) -> pool 

56 self._pools: Dict[tuple, PoolType] = {} 

57 # Weak references to event loops for cleanup 

58 self._loop_refs: WeakValueDictionary = WeakValueDictionary() 

59 # Register cleanup on exit 

60 atexit.register(self._cleanup_on_exit) 

61 

62 async def get_pool( 

63 self, 

64 config: BasePoolConfig, 

65 create_pool_func: Callable[[BasePoolConfig], Awaitable[PoolType]], 

66 validate_pool_func: Optional[Callable[[PoolType], Awaitable[None]]] = None, 

67 close_pool_func: Optional[Callable[[PoolType], Awaitable[None]]] = None 

68 ) -> PoolType: 

69 """ 

70 Get or create a connection pool for the current event loop. 

71  

72 Args: 

73 config: Pool configuration 

74 create_pool_func: Async function to create a new pool 

75 validate_pool_func: Optional async function to validate an existing pool 

76 close_pool_func: Optional async function to close a pool 

77  

78 Returns: 

79 Pool instance for the current event loop 

80 """ 

81 loop = asyncio.get_running_loop() 

82 loop_id = id(loop) 

83 config_hash = hash(config.to_hash_key()) 

84 pool_key = (config_hash, loop_id) 

85 

86 # Check if we already have a pool for this config and loop 

87 if pool_key in self._pools: 

88 pool_entry = self._pools[pool_key] 

89 # Handle both old and new format 

90 if isinstance(pool_entry, tuple): 

91 pool, _ = pool_entry 

92 else: 

93 pool = pool_entry 

94 

95 # Validate the pool if validation function provided 

96 if validate_pool_func: 

97 try: 

98 await validate_pool_func(pool) 

99 return pool 

100 except Exception as e: 

101 logger.warning(f"Pool for loop {loop_id} is invalid: {e}. Creating new one.") 

102 await self._close_pool(pool_key, close_pool_func) 

103 else: 

104 return pool 

105 

106 # Create new pool 

107 logger.info(f"Creating new connection pool for loop {loop_id}") 

108 pool = await create_pool_func(config) 

109 

110 # Store pool and loop reference with close function 

111 self._pools[pool_key] = (pool, close_pool_func) 

112 self._loop_refs[loop_id] = loop 

113 

114 return pool 

115 

116 async def _close_pool(self, pool_key: tuple, close_func: Optional[Callable] = None): 

117 """Close and remove a pool.""" 

118 if pool_key in self._pools: 

119 pool_entry = self._pools[pool_key] 

120 # Handle both old format (pool) and new format (pool, close_func) 

121 if isinstance(pool_entry, tuple): 

122 pool, stored_close_func = pool_entry 

123 close_func = close_func or stored_close_func 

124 else: 

125 pool = pool_entry 

126 

127 try: 

128 # Check if we have a running event loop 

129 try: 

130 loop = asyncio.get_running_loop() 

131 if loop.is_closed(): 

132 # Event loop is closed, skip async cleanup 

133 return 

134 except RuntimeError: 

135 # No running event loop, skip async cleanup 

136 return 

137 

138 if close_func: 

139 await close_func(pool) 

140 elif hasattr(pool, 'close'): 

141 await pool.close() 

142 except RuntimeError as e: 

143 # Silently ignore "Event loop is closed" errors 

144 if "Event loop is closed" not in str(e): 

145 logger.error(f"Error closing pool: {e}") 

146 except Exception as e: 

147 logger.error(f"Error closing pool: {e}") 

148 finally: 

149 del self._pools[pool_key] 

150 

151 async def remove_pool(self, config: BasePoolConfig) -> bool: 

152 """ 

153 Remove a pool for the current event loop. 

154  

155 Args: 

156 config: Pool configuration 

157  

158 Returns: 

159 True if pool was removed, False if not found 

160 """ 

161 loop_id = id(asyncio.get_running_loop()) 

162 config_hash = hash(config.to_hash_key()) 

163 pool_key = (config_hash, loop_id) 

164 

165 if pool_key in self._pools: 

166 await self._close_pool(pool_key) 

167 return True 

168 return False 

169 

170 async def close_all(self): 

171 """Close all connection pools.""" 

172 for pool_key in list(self._pools.keys()): 

173 await self._close_pool(pool_key) 

174 

175 def get_pool_count(self) -> int: 

176 """Get the number of active pools.""" 

177 return len(self._pools) 

178 

179 def get_pool_info(self) -> Dict[str, Any]: 

180 """Get information about all active pools.""" 

181 info = {} 

182 for (config_hash, loop_id), pool_entry in self._pools.items(): 

183 # Handle both old and new format 

184 if isinstance(pool_entry, tuple): 

185 pool, _ = pool_entry 

186 else: 

187 pool = pool_entry 

188 

189 key = f"config_{config_hash}_loop_{loop_id}" 

190 info[key] = { 

191 "loop_id": loop_id, 

192 "config_hash": config_hash, 

193 "pool": str(pool) 

194 } 

195 return info 

196 

197 def _cleanup_on_exit(self): 

198 """Cleanup function called on program exit.""" 

199 if self._pools: 

200 logger.debug(f"Cleaning up {len(self._pools)} connection pools on exit") 

201 # Try to get any running loop 

202 try: 

203 loop = asyncio.get_running_loop() 

204 except RuntimeError: 

205 # No running loop, try to create one 

206 loop = asyncio.new_event_loop() 

207 asyncio.set_event_loop(loop) 

208 try: 

209 loop.run_until_complete(self.close_all()) 

210 finally: 

211 loop.close() 

212 else: 

213 # There's a running loop, schedule cleanup 

214 asyncio.create_task(self.close_all())