Coverage for streamwise/http_session_manager.py: 100%

1""" 

2HTTP session manager for aiohttp ClientSession. 

3""" 

4 

5import logging 

6 

7from aiohttp import ClientSession 

8from aiohttp import ClientTimeout 

9from aiohttp import TCPConnector 

10 

11 

12client_session: ClientSession | None = None 

13 

14# URL scheme used when constructing outbound service URLs (http or https). 

15# Call set_service_scheme("https") at startup to enable HTTPS for all 

16# service-to-service communication (health checks, file fetches, job 

17# submissions, etc.). 

18SERVICE_SCHEME: str = "http" 

19 

20# Whether to verify SSL certificates when making outbound HTTPS requests. 

21# Set to False when services use self-signed certificates (e.g. Key Vault 

22# generated certs in AKS). Call set_verify_ssl(False) at startup alongside 

23# set_service_scheme("https") to disable certificate verification. 

24VERIFY_SSL: bool = True 

25 

26 

27def set_service_scheme(scheme: str) -> None: 

28 """Set the URL scheme used for outbound service requests.""" 

29 global SERVICE_SCHEME 

30 if scheme not in ("http", "https"): 

31 raise ValueError(f"Invalid service scheme: {scheme!r}. Must be 'http' or 'https'.") 

32 SERVICE_SCHEME = scheme 

33 

34 

35def set_verify_ssl(verify: bool) -> None: 

36 """Set whether to verify SSL certificates for outbound HTTPS requests.""" 

37 global VERIFY_SSL 

38 VERIFY_SSL = verify 

39 

40 

41def create_client_session_instance() -> ClientSession: 

42 connector = TCPConnector( 

43 limit=100, 

44 limit_per_host=10, 

45 use_dns_cache=True, 

46 ssl=VERIFY_SSL) 

47 timeout = ClientTimeout( 

48 total=0.5, 

49 connect=0.5) 

50 return ClientSession( 

51 connector=connector, 

52 timeout=timeout) 

53 

54 

55async def startup() -> None: 

56 """Initialize sessions before the server starts.""" 

57 global client_session 

58 if client_session is None or client_session.closed: 

59 logging.info("Creating aiohttp client session...") 

60 client_session = create_client_session_instance() 

61 

62 

63async def shutdown() -> None: 

64 """Cleanup tasks after server stops.""" 

65 global client_session 

66 if client_session and not client_session.closed: 

67 logging.info("Closing aiohttp client session...") 

68 await client_session.close() 

69 client_session = None 

70 

71 

72async def get_global_session() -> ClientSession: 

73 """Get or create a global aiohttp ClientSession for reuse.""" 

74 global client_session 

75 if client_session is None or client_session.closed: 

76 client_session = create_client_session_instance() 

77 return client_session