Coverage for model/utils/request.py: 88%

1import inspect 

2import httpx 

3from fastapi import Depends, Request 

4from mongo import engine 

5from .response import NOJException 

6 

7__all__ = ('get_doc', 'get_ip', 'get_http_client') 

8 

9 

10def get_doc(src_param: str, cls, param_type=str): 

11 """ 

12 Depends factory: resolves a MongoEngine document from a path/query parameter. 

13 

14 Usage: 

15 @router.get('/{username}') 

16 def handler(target_user = get_doc('username', User)): ... 

17 

18 @router.get('/{problem_id}') 

19 def handler(problem = get_doc('problem_id', Problem, int)): ... 

20 """ 

21 

22 def dependency(**kwargs): 

23 val = kwargs[src_param] 

24 try: 

25 doc = cls(val) 

26 if not doc: 

27 raise engine.DoesNotExist(f'{cls.__name__} not found') 

28 return doc 

29 except engine.DoesNotExist as e: 

30 raise NOJException(str(e), 404) 

31 except engine.ValidationError: 

32 raise NOJException('Invalid parameter', 400) 

33 

34 # Give the function a proper named signature so FastAPI can inject the path param 

35 param = inspect.Parameter( 

36 src_param, 

37 kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, 

38 annotation=param_type, 

39 ) 

40 dependency.__signature__ = inspect.Signature([param]) 

41 dependency.__annotations__ = {src_param: param_type} 

42 return Depends(dependency) 

43 

44 

45def get_http_client(request: Request) -> httpx.Client: 

46 return request.app.state.http_client 

47 

48 

49def get_ip(request: Request) -> str: 

50 # cf-connecting-ip is set by Cloudflare and is more reliable than X-Forwarded-For 

51 cf_ip = request.headers.get('cf-connecting-ip', '').strip() 

52 if cf_ip: 

53 return cf_ip 

54 forwarded_for = request.headers.get('X-Forwarded-For', 

55 '').split(',')[-1].strip() 

56 if forwarded_for: 

57 return forwarded_for 

58 return request.client.host if request.client else ''